Ml Matgl Finetune

MatGL Fine-tuning

Goal

To evaluate and improve the accuracy of a foundation MatGL potential (e.g., CHGNet, M3GNet, TensorNet) for a specific chemical system or physical property using the provided Python fine-tuning script.

Instructions

1. Prepare Labeled Dataset: Obtain diverse structures with high-fidelity labels (energy, forces, stress). See the /benchmark-finetuning workflow for details.
2. Custom Data Conversion: Read the source data format and write a customized conversion script if needed, formatting it for the subsequent preparation step.
3. Data Preparation: Execute scripts/prepare_matgl_data.py to process JSON structures and split into training and validation sets.
4. Fine-Tuning: Execute scripts/train_matgl.py to begin fine-tuning natively on the GPU using PyTorch Lightning.
5. Validation: Verify convergence and compare against the benchmarked foundation metrics.
6. Registration: Use the register_model tool to register the newly fine-tuned model checkpoint into the local registry so future research tasks can discover and reuse it.

Usage

1. Data Preparation

Convert your dataset into the appropriate JSON format for MatGL training:

conda run -n matgl-agent python .agents/skills/ml-matgl-finetune/scripts/prepare_matgl_data.py \
    --data /path/to/training_data.json \
    --model CHGNet-MatPES-PBE-2025.2.10-2.7M-PES \
    --val-split 0.1 \
    --output-dir ./matgl_finetuned

2. Run Training

Fine-tune the model using the prepared data:

conda run -n matgl-agent python .agents/skills/ml-matgl-finetune/scripts/train_matgl.py \
    --train-data ./matgl_finetuned/train_data.json \
    --val-data ./matgl_finetuned/val_data.json \
    --model CHGNet-MatPES-PBE-2025.2.10-2.7M-PES \
    --epochs 10 \
    --lr 1e-3 \
    --batch-size 4 \
    --freeze-backbone \
    --output-dir ./matgl_finetuned

Training Configuration

MatGL fine-tuning is divided into a data preparation step (formatting nested dictionaries and converting lists) and a native training run utilizing PyTorch Lightning.

Data Preparation Arguments (prepare_matgl_data.py)

Training Arguments (train_matgl.py)

Constraints

• Data Size: For small datasets, --freeze-backbone is strongly recommended to prevent catastrophic forgetting.
• Reference Energies (element_refs): If your fine-tuning data is computed using the same DFT functional (e.g., PBE) as the foundation model's original training data, you should reuse the foundation model's original isolated atom reference energies instead of re-fitting them. This maintains thermodynamic compatibility across the periodic table.
• Environment: Must be executed within the matgl-agent conda environment where MatGL and DGL are properly configured.
• Stress Units: MatGL inherently converts stress internally to GPa, however the standard expected inputs directly into its JSON files are eV/ų. Raw VASP stress obtained directly via some JSON files may be in kilo-Bar (kB). The Atomate2 MCP tool handles this conversion automatically when convert_units=True. However, if your JSON labels contain raw kB stress, you MUST pass the --vasp-stress-conversion flag to scripts/prepare_matgl_data.py to automatically scale them by -1/160.2x. For more details on unit standardization, see @[.agents/skills/general-property-units/SKILL.md].

Author: Bowen Deng Contact: GitHub @learningmatter-mit