learningmatter-mit Fairchem Fine-tuning
Goal
To evaluate and improve the accuracy of a foundation Fairchem potential (e.g., UMA, ESEN) for a specific chemical system or physical property using the provided Python fine-tuning script.
Instructions
- Prepare Labeled Dataset: Obtain diverse structures with high-fidelity labels (energy, forces, stress). See the
/benchmark-finetuningworkflow for details. - Custom Data Conversion: Read the source data format and write a customized conversion script if needed, formatting it for the subsequent preparation step.
- Benchmarking: Predict results on the new labels and benchmark the foundation model using ml-mlip-benchmark.
- Data Preparation: Execute
scripts/prepare_fairchem_data.pyto convert JSON structures to extxyz, generate native LMDB databases, compute dataset references, and configure a templateduma_sm_finetune_template.yaml. - Fine-Tuning: Execute
fairchem -c uma_sm_finetune_template.yaml job.run_dir=XXXnatively. - Validation: Run
scripts/extract_fairchem_logs.pyto extract curves and verify convergence against the benchmarked foundation metrics. - Registration: Use the
register_modeltool to register the newly fine-tuned model checkpoint into the local registry so future research tasks can discover and reuse it.
Training Configuration
Fairchem fine-tuning relies heavily on the fairchem CLI, which uses Hydra for configuration. The script scripts/prepare_fairchem_data.py bridges standard data into the complex Fairchem directory structure and generates .aselmdb dataset formats automatically.
Basic Arguments (Data Prep Script)
| Key | Type | Default | Description |
|---|---|---|---|
--data |
str | (Required) | Path to JSON file containing ASE/pymatgen structure dictionaries |
--val-data |
str | None | Path to JSON file containing validation split. (Optional, otherwise --val-split is used) |
--val-split |
float | 0.1 | Validation split if --val-data is not provided |
--seed |
int | 42 | Random seed for data splitting and initialization |
--model |
str | uma-s-1p1 |
Base model name or path to a checkpoint |
--task-name |
str | omat |
The specific multi-task context to run against (omat, omol) |
--epochs |
int | 10 | Number of training epochs |
--lr |
float | 4e-4 | Peak learning rate for training |
--batch-size |
int | 2 | Training batch size |
--freeze-backbone |
flag | N/A | Add flag to mathematically freeze OCP/UMA interaction layers |
--weight-decay |
float | 1e-3 | Weight decay parameter |
--warmup-factor |
float | 0.2 | LR warmup factor |
--warmup-epochs |
float | 0.01 | Epochs to perform LR warmup |
--lr-min-factor |
float | 0.01 | Minimum LR factor after decay |
--clip-grad-norm |
float | 100.0 | Gradient clipping threshold |
--evaluate-every-n-steps |
int | 100 | Steps frequency for validation evaluation |
--checkpoint-every-n-steps |
int | 1000 | Steps frequency for model checkpointing |
--ema-decay |
float | 0.999 | Exponential moving average decay parameter |
--linref-coeff |
str | None | JSON array of elemental energy linear references. If None, it auto-computes it over the data. |
--vasp-stress-conversion |
flag | N/A | Add flag to automatically convert kB to eV/ų for VASP inputs |
--output-dir |
str | ./fairchem_finetuning |
Directory to save the lmdb_output intermediate data and run configs |
[!NOTE] If you have created a dedicated research directory for your current workflow (e.g. using the
create_research_dirtool), you should set the--output-dirargument to a folder within that active research directory to keep all artifacts and models organized. The data preparation script takes several minutes because it automatically creates.aselmdbcopies of all structural inputs mapping to specific index structures.
[!WARNING] Stress Units: Fairchem 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 whenconvert_units=True. However, if your JSON labels contain rawkBstress, you MUST pass the--vasp-stress-conversionflag toscripts/prepare_fairchem_data.pyto automatically scale them by-1/160.2x. For more details on unit standardization, see @[.agents/skills/general-property-units/SKILL.md].
Usage:
# 1. Prepare Data and Config
conda run -n fairchem-agent python .agents/skills/ml-fairchem-finetune/scripts/prepare_fairchem_data.py \
--data /path/to/training_data.json \
--model uma-s-1p1 \
--epochs 10 \
--lr 4e-4 \
--batch-size 2 \
--freeze-backbone \
--output-dir ./research/my_dir/fairchem_finetuning
# 2. Run Training
export PYTHONPATH=/path/to/research/my_dir/fairchem_finetuning/lmdb_output:$PYTHONPATH
cd /path/to/research/my_dir/fairchem_finetuning/lmdb_output
conda run -n fairchem-agent fairchem -c uma_sm_finetune_template.yaml job.run_dir=/path/to/research/my_dir/fairchem_finetuning/runs +job.timestamp_id=run_10ep
# 3. Extract Training Logs (Optional, to create standard training_history.json)
conda run -n fairchem-agent python .agents/skills/ml-fairchem-finetune/scripts/extract_fairchem_logs.py \
--log /path/to/research/my_dir/fairchem_finetuning/runs/run_10ep/logs/trainer.log \
--output-dir /path/to/research/my_dir/fairchem_finetuning/results
Constraints
- Multi-task Setup: UMA and ESEN are trained explicitly on tasks. Be absolutely sure to specify the right
--task-namefor your datasets (omatvsomol). - Reference Energies (
linref_coeff): 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 extract and pass the foundation model's originallinref_coeffarray instead of allowing the script to automatically re-fit it via Least Squares. This maintains thermodynamic scale compatibility across the periodic table. - GPU Overhead: Fairchem configuration files compile PyTorch networks prior to run and memory overhead can cause execution to take over 5 minutes to generate logs if using data parallel or multi-gpu execution. Disable wandb using the
--debugparameter within the configuration file if training stalls completely.
Author: Bowen Deng Contact: GitHub @learningmatter-mit