Fine-Tuning Universal Models
Overview
The fine-tuning module allows you to adapt a pre-trained universal ML potential to your specific chemical system using a small amount of DFT reference data. This is an alternative to training from scratch and typically requires far fewer DFT calculations to reach the same accuracy.
Currently supported backend:
DeePMD — fine-tunes DPA-1, DPA-2, or DPA-3 models via
dp --pt finetune(requires DeePMD-kit v3 with PyTorch backend)
SPARC trains num_models independent models (one per training_x/
directory). These are used as a committee for Query-by-Committee uncertainty
estimation during active learning.
Note
Fine-tuning replaces the standard MLIP training step. Set finetune.enabled: true
in input.yaml and ensure mlip_setup.training: true is also enabled.
Configuration
Enable fine-tuning in input.yaml under the finetune section:
finetune:
enabled: true
model_type: "deepmd"
pretrained_model: "DPA3.pt" # Path to pre-trained model file
model_branch: "Omat24" # Multi-task branch for DPA-3 (optional)
input_file: null # Fine-tune JSON (uses mlip_setup.input_file if null)
learning_rate: 0.001 # Starting learning rate
device: "cpu" # "cpu" or "cuda"
All fields except enabled and model_type are optional and have sensible defaults.
DeePMD Fine-Tuning
Uses dp --pt train --finetune to initialise from a pre-trained checkpoint.
Requirements
DeePMD-kit v3 with PyTorch backend (
pip install deepmd-kit[torch])A pre-trained model file (
*.pt) — e.g., DPA-3 from the DeepModeling model hub
What it does
For each model i in 1 … num_models:
Copies and updates
input.json(data paths, atom types, optional learning rate override)Runs
dp --pt train input.json --finetune <pretrained_model.pt> [--model-branch BRANCH]Freezes the trained model:
dp --pt freeze -o frozen_model_i.pth
The frozen models are saved inside 01.train/training_i/frozen_model_i.pth.
Multi-task models (DPA-3)
DPA-3 is a multi-task model trained on several datasets simultaneously. Use
model_branch to select the branch that best matches your system:
finetune:
model_type: "deepmd"
pretrained_model: "DPA3.pt"
model_branch: "Omat24" # inorganic materials
# model_branch: "Organic_Reactions"
Leave model_branch as null for single-task models (DPA-1, DPA-2).
Integration with Active Learning
When fine-tuning is enabled, SPARC substitutes it for the standard DeePMD training step inside the active learning loop:
Iteration N
├── 00.dft/ DFT single-point labelling
├── 01.train/ Fine-tuning (replaces from-scratch training)
│ ├── training_1/
│ │ ├── input.json
│ │ └── frozen_model_1.pth
│ ├── training_2/
│ └── ...
└── 02.dpmd/ ML-MD + model deviation (QbC)
The frozen models produced here are picked up automatically by the ML-MD and model-deviation steps.
Troubleshooting
- “Pre-trained model not found”
Ensure
pretrained_modelis a valid path relative to the project root.- “DeePMD fine-tuning requires v3 or later”
Upgrade with
pip install --upgrade deepmd-kit[torch].- “unexpected keyword argument” during fine-tuning
Version mismatch between the pre-trained model and installed DeePMD-kit. Either upgrade DeePMD-kit or use a model compatible with your installed version (e.g., DPA-2 models work with DeePMD-kit v3.0.x).
Module Contents
Fine-tuning module for universal ML potentials.
Supports:
DeePMD fine-tuning via DeePMD-kit v3 (dp –pt finetune) Works with DPA-1, DPA-2, DPA-3 and other DeePMD models
MACE fine-tuning via mace_run_train –foundation_model
Both produce frozen models compatible with the SPARC active learning workflow.
- sparc.src.finetune.deepmd_finetune(datadir, atom_types, training_dir, num_models, input_file='input.json', pretrained_model='DPA3.pt', model_branch=None, learning_rate=None, **kwargs)[source]
Fine-tune a DeePMD universal model on system-specific DFT data.
Uses DeePMD-kit v3 dp –pt finetune command which initializes from a pre-trained checkpoint (DPA-1, DPA-2, DPA-3, etc.) and adapts to local data.
- Parameters:
datadir (
str) – Path to directory containing training_data/ and validation_data/atom_types (
List[str]) – List of atomic species (e.g., [‘O’, ‘H’])training_dir (
str) – Directory where training_1/, training_2/, etc. will be creatednum_models (
int) – Number of models to fine-tune (for ensemble/QbC)input_file (
str) – Path to DeepMD input JSON configurationpretrained_model (
str) – Path to pre-trained model file (.pt)learning_rate (
Optional[float]) – Override starting learning rate for fine-tuningmodel_branch (str | None)
- Returns:
Name of the frozen model file (e.g., ‘frozen_model_2.pth’)
- Return type:
- sparc.src.finetune.finetune_training(finetune_config, datadir, atom_types, training_dir, num_models, input_file='input.json')[source]
Dispatch fine-tuning to the appropriate backend (DeePMD or MACE).
- Parameters:
- Returns:
Frozen model filename
- Return type:
- sparc.src.finetune.mace_finetune(datadir, atom_types, training_dir, num_models, pretrained_model='medium', num_epochs=100, learning_rate=0.001, batch_size=4, device='cpu', stress_key='stress', **kwargs)[source]
Fine-tune MACE foundation model on system-specific DFT data.
Uses mace_run_train with –foundation_model flag for transfer learning from pre-trained MACE-MP-0 models.
- Parameters:
datadir (
str) – Path to directory containing training_data/ and validation_data/atom_types (
List[str]) – List of atomic species (e.g., [‘O’, ‘H’])training_dir (
str) – Directory where training_1/, training_2/, etc. will be creatednum_models (
int) – Number of models to fine-tune (for ensemble/QbC)pretrained_model (
str) – Foundation model name (‘small’, ‘medium’, ‘large’) or path to .model filenum_epochs (
int) – Number of fine-tuning epochs (default: 100)learning_rate (
float) – Learning rate for fine-tuning (default: 0.001)batch_size (
int) – Batch size (default: 4)device (
str) – Device for training: ‘cpu’ or ‘cuda’ (default: ‘cpu’)stress_key (str)
- Returns:
Name of the fine-tuned model file
- Return type: