Ml Mace Finetune

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.
4. Data Preparation: Execute scripts/prepare_mace_data.py to convert JSON structures to .xyz data files.
5. Config Generation: Execute scripts/generate_mace_config.py using the .xyz data to produce finetune_config.yaml.
6. Fine-Tuning: Execute mace_run_train --config /path/to/finetune_config.yaml to begin fine-tuning natively on the GPU.
7. Validation: Verify convergence and compare against the benchmarked foundation metrics.
8. 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)

Basic Arguments (Configuration Generation Script)

[!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)

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

LR Scheduler

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

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

Loss Function

[!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 @[.agents/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 — MACE-OMAT-0-small, 10 epochs

Usage:

# 1. Prepare Data
conda run -n mace-agent python .agents/skills/ml-mace-finetune/scripts/prepare_mace_data.py \
    --data /path/to/training_data.json \
    --output-dir ./mace_finetuned_data

# 2. Generate Configuration
conda run -n mace-agent python .agents/skills/ml-mace-finetune/scripts/generate_mace_config.py \
    --train-file ./mace_finetuned_data/train.xyz \
    --valid-file ./mace_finetuned_data/valid.xyz \
    --model MACE-OMAT-0-small \
    --epochs 10 \
    --lr 1e-4 \
    --batch-size 2 \
    --freeze-backbone \
    --output-dir ./mace_finetuned

# 3. Run Training
conda run -n mace-agent mace_run_train --config ./mace_finetuned/finetune_config.yaml

# 4. Extract Training Logs (Optional, to create standard training_history.json)
conda run -n base-agent python .agents/skills/ml-mace-finetune/scripts/extract_mace_logs.py \
    --results-dir ./mace_finetuned/results

Examples

See the mace-wbm-finetune directory for a complete, runnable example of 10-epoch fine-tuning on high-energy crystal structures (WBM dataset), including exact usage of the --vasp-stress-conversion flag and diagnostic output artifacts.

Constraints

• Data Size: For small datasets (<500 structures), freeze_backbone=True is strongly recommended.
• Reference Energies (E0s): 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 (E0s) instead of re-fitting them. This maintains thermodynamic compatibility across the periodic table for elements not in your fine-tuning set.
• Units (input): Stress labels must be in eV/ų as per project standards.
• Units (output): All training_history.json files use meV units: energy MAE in meV/atom, force MAE in meV/Å, stress MAE in meV/ų.

Author: Bowen Deng Contact: GitHub @learningmatter-mit