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.