ml-matgl-finetune

Files (10)

• SKILL.md
• README.md · examples/matgl-wbm-finetune/
• checkpoint_roundtrip.json · examples/matgl-wbm-finetune/
• finetune_record.json · examples/matgl-wbm-finetune/
• training_history.json · examples/matgl-wbm-finetune/
• training_history.png · examples/matgl-wbm-finetune/
• verify_checkpoint.py · examples/matgl-wbm-finetune/
• generate_matgl_config.py · scripts/
• prepare_matgl_data.py · scripts/
• train_matgl.py · scripts/

SKILL.md (first lines)

# MatGL Fine-tuning

## Goal
To evaluate and improve the accuracy of a foundation MatGL potential (e.g., CHGNet, 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:
```bash
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/prepare_matgl_data.py \
    --data /path/to/training_data.json \
    --model CHGNet-PES-MatPES-PBE-1M-2026.9 \
    --val-split 0.1 \
    --output-dir ./matgl_finetuned
```

### 2. Run Training
Fine-tune the model using the prepared data:
```bash
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/train_matgl.py \
    --train-data ./matgl_finetuned/train_data.json \
    --val-data ./matgl_finetuned/val_data.json \
    --model CHGNet-PES-MatPES-PBE-1M-2026.9 \
    --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`)

| Key | Type | Default | Description |
|:----|:-----|:--------|:------------|
| `--data` | str | (Required) | Path to JSON file containing ASE/pymatgen structure dictionaries |
| `--model` | str | `CHGNet-PES-MatPES-PBE-1M-2026.9` | Base model name or path to a checkpoint |
| `--output-dir` | str | `./fine_tuning` | Directory to save the processed data |
| `--val-split` | float | 0.1 | Fraction of data to set aside for validation |
| `--seed` | int | 42 | Random seed for splitting validation data |
| `--vasp-stress-conversion`| flag | - | If set, multiplies stress values by -1/160.2x to convert VASP raw kB to eV/ų |

### Training Arguments (`train_matgl.py`)

| Key | Type | Default | Description |
|:----|:-----|:--------|:------------|
| `--train-data` | str | (Required) | Path to JSON file containing training data |
| `--val-data` | str | None | Path to JSON file containing validation data (optional) |
| `--model` | str | `CHGNet-PES-MatPES-PBE-1M-2026.9` | Base model name or path to a checkpoint |
| `--epochs` | int | 10 | Number of training epochs |
| `--lr` | float | 1e-3 | Learning rate |
| `--batch-size` | int | 4 | Training batch size |
| `--device` | str | `auto` | Target compute device (`cuda` or `cpu`) |
| `--output-dir` | str | `./fine_tuning` | Directory to save the fine-tuned model and logs |
| `--freeze-backbone` | flag | - | Freeze backbone (interaction blocks); only readout heads are trained |
| `--reinit-head` | flag | - | Re-initialize readout head weights |
| `--scheduler` | str | `CosineAnnealingLR` | `CosineAnnealingLR` or `ReduceLROnPlateau` |
| `--patience` | int | None | Early stopping patience (epochs) |
| `--energy-weight` | float | 1.0 | Loss weight for energy |
| `--force-weight` | float | 1.0 | Loss weight for forces |
| `--stress-weight` | float | 0.1 | Loss weight for stress |

## 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**: Runs in the `mlip` environment (`venv/run mlip ...`). MatGL >= 4 uses PyTorch Geometric only; there is no DGL backend.
- **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 @[skills/general-property-units/SKILL.md].

---

**Author:** Bowen Deng
**Contact:** [GitHub @learningmatter-mit](https://github.com/learningmatter-mit)

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills