ml-mace-finetune

Files (10)

• SKILL.md
• README.md · examples/mace-wbm-finetune/
• finetune_config.yaml · examples/mace-wbm-finetune/
• label_distributions.png · examples/mace-wbm-finetune/
• training_history.json · examples/mace-wbm-finetune/
• training_history.png · examples/mace-wbm-finetune/
• valid.xyz · examples/mace-wbm-finetune/
• extract_mace_logs.py · scripts/
• generate_mace_config.py · scripts/
• prepare_mace_data.py · scripts/

SKILL.md (first lines)

# MACE Fine-tuning

## Goal
To evaluate and improve the accuracy of a foundation MACE potential for a specific chemical system or physical property using the provided Python fine-tuning script and data-augmentation.

## 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.  **Benchmarking**: Predict results on the new labels and benchmark the foundation model using [ml-mlip-benchmark](../ml-mlip-benchmark/SKILL.md).
16. **Data Preparation**: Execute `scripts/prepare_mace_data.py` to convert JSON structures to `.xyz` data files.
17. **Config Generation**: Execute `scripts/generate_mace_config.py` using the `.xyz` data to produce `finetune_config.yaml`.
18. **Fine-Tuning**: Execute `mace_run_train --config /path/to/finetune_config.yaml` to begin fine-tuning natively on the GPU.
19. **Validation**: Verify convergence and compare against the benchmarked foundation metrics.
20. **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.

## Training Configuration

MACE fine-tuning is divided into a data preparation step, a configuration generation step, and a standard native training run. The script `scripts/prepare_mace_data.py` generates `.xyz` files, and `scripts/generate_mace_config.py` converts arguments into a fully-formed `finetune_config.yaml` configuration compatible with the MACE default parser.

### Basic Arguments (Data Prep Script)

| Key | Type | Default | Description |
|:----|:-----|:--------|:------------|
| `--data` | str | (Required) | Path to JSON file containing ASE/pymatgen structure dictionaries |
| `--output-dir` | str | `./fine_tuning_data` | Directory to save the converted .xyz data |
| `--val-split` | float | 0.1 | Fraction of data to set aside for validation |
| `--seed` | int | 42 | Random seed for validation splitting |
| `--vasp-stress-conversion`| flag | - | If set, multiplies stress values by -1/160.2x to convert VASP raw kB to eV/ų |

### Basic Arguments (Configuration Generation Script)

| Key | Type | Default | Description |
|:----|:-----|:--------|:------------|
| `--train-file` | str | (Required) | Path to the converted `train.xyz` data |
| `--valid-file` | str | None | Path to the `valid.xyz` data |
| `--model` | str | `MACE-MP-small` | Base model name or path to a checkpoint |
| `--epochs` | int | 100 | Number of training epochs |
| `--lr` | float | 0.01 | Peak learning rate for training |
| `--batch-size` | int | 2 | Training batch size |
| `--output-dir` | str | `./fine_tuning` | Directory to save the fine-tuned model and logs |
| `--device` | str | `cuda` | Target compute device (`cuda` or `cpu`) |

> [!NOTE]
> If you have created a dedicated research directory for your current workflow (e.g. using the `create_research_dir` tool), you should set the `--output-dir` argument to a folder within that active research directory to keep all artifacts and models organized.

### Model Freezing and Heads (Configuration Generation Script)

| Key | Type | Default | Choices | Description |
|:----|:-----|:--------|:--------|:------------|
| `--freeze-backbone` | flag | N/A | Add flag | Freeze backbone (interaction blocks); only readout heads are trained. |
| `--reinit-head` | flag | N/A | Add flag | Re-initialize readout weights. When absent (default), pre-trained readout is preserved. |
| `--multiheads` | flag | N/A | Add flag | Enable multi-head fine-tuning (adds new head while keeping existing ones). |

### Advanced Parameters (Manual YAML Injection)

> [!IMPORTANT]
> The following parameters govern Optimizer, Regularization, and Scheduling. They are **NOT** exposed via the `prepare_mace_data.py` CLI. Note that `prepare_mace_data.py` exposes `--energy-weight`, `--forces-weight` (default 10.0), and `--stress-weight` which inject cleanly into the initial YAML. For all other properties below, you must manually append the keys to the generated `finetune_config.yaml` file prior to running `mace_run_train`.

#### Optimizer & Regularization

| Key | Type | Default | Choices / Range | Description |
|:----|:-----|:--------|:----------------|:------------|
| `optimizer` | str | `"adam"` | `"adam"`, `"adamw"`, `"schedulefree"` | Optimizer type. |
| `weight_decay` | float | 5e-7 | ≥0 | L2 weight decay. |
| `amsgrad` | bool | True | `True`, `False` | Use AMSGrad variant of Adam. |
| `clip_grad` | float | 10.0 | >0 or None | Maximum gradient norm for clipping. Set to None to disable. |
| `ema` | bool | True | `True`, `False` | Enable exponential moving average of model weights. |
| `ema_decay` | float | 0.99 | 0–1 | EMA decay rate. Higher = more smoothing. |

### LR Scheduler

| Key | Type | Default | Choices / Range | Description |
|:----|:-----|:--------|:----------------|:------------|
| `scheduler` | str | `"ReduceLROnPlateau"` | `"ReduceLROnPlateau"`, `"ExponentialLR"` | LR scheduler type. |
| `lr_factor` | float | 0.8 | 0–1 | Factor by which LR is reduced on plateau (for ReduceLROnPlateau). |
| `scheduler_patience` | int | 50 | ≥1 | Epochs without improvement before reducing LR (for ReduceLROnPlateau). |
| `lr_scheduler_gamma` | float | 0.9993 | 0–1 | Per-epoch multiplicative decay factor (for ExponentialLR). |

**ReduceLROnPlateau** (default): Monitors validation loss. When loss stops improving for `scheduler_patience` epochs, LR is multiplied by `lr_factor`.

**ExponentialLR**: LR decays every epoch by `lr_scheduler_gamma`. At epoch *n*: `LR = learning_rate × lr_scheduler_gamma^n`.

### Early Stopping

| Key | Type | Default | Choices / Range | Description |
|:----|:-----|:--------|:----------------|:------------|
| `patience` | int | 2048 | ≥1 | Stop training after this many epochs without improvement. |

> [!NOTE]
> Default `patience=2048` effectively disables early stopping. Set lower (e.g., 100–200) if you want to stop early.

### Loss Function

| Key | Type | Default | Choices | Description |
|:----|:-----|:--------|:--------|:------------|
| `loss` | str | `"weighted"` | `"weighted"`, `"universal"`, etc. | Loss function type. `"universal"` is auto-set by data script when stress data is detected. |
| `energy_weight` | float | 1.0 | ≥0 | Weight for energy loss (exposed via CLI). |
| `forces_weight` | float | 10.0 | ≥0 | Weight for forces loss (exposed via CLI). |
| `stress_weight` | float | 1.0 | ≥0 | Weight for stress loss (exposed via CLI). |
| `compute_forces` | bool | True | `True`, `False` | Include forces in training. |
| `compute_stress` | bool | False | `True`, `False` | Include stress in training (auto-enabled by data script). |

> [!WARNING]
> **Stress Units**: MACE expects stress in `eV/ų`. Raw VASP stress obtained directly via some JSON files may be in kilo-Bar (`kB`), which is ~160x larger and will cause catastrophic training divergence. 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_mace_data.py` to automatically scale them by `-1/160.2x`. For more details on unit standardization, see @[skills/general-property-units/SKILL.md].

> [!WARNING]
> **Learning rate sensitivity for MACE-OMAT**: The official MACE docs recommend `lr=0.01` for MACE-MP-0, but MACE-OMAT-0-small requires `lr=1e-4` to avoid divergence. Higher values (1e-3, 0.01) cause catastrophic forgetting even with frozen backbone + EMA.

> [!IMPORTANT]
> The generated `finetune_config.yaml` maps exactly to `mace_run_train` arguments. You can open and manually modify the YAML file before running `mace_run_train` to inject any advanced parameter.

- [MACE Si-O fine-tuning with frozen backbone](examples/mace-sio-frozen/) — MACE-OMAT-0-small, 10 epochs

Usage:
```bash
# 1. Prepare Data
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/prepare_mace_data.py \
    --data /path/to/training_data.json \
    --output-dir ./mace_finetuned_data

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills