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:

  1. Copies and updates input.json (data paths, atom types, optional learning rate override)

  2. Runs dp --pt train input.json --finetune <pretrained_model.pt> [--model-branch BRANCH]

  3. 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_model is 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:

  1. DeePMD fine-tuning via DeePMD-kit v3 (dp –pt finetune) Works with DPA-1, DPA-2, DPA-3 and other DeePMD models

  2. 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 created

  • num_models (int) – Number of models to fine-tune (for ensemble/QbC)

  • input_file (str) – Path to DeepMD input JSON configuration

  • pretrained_model (str) – Path to pre-trained model file (.pt)

  • learning_rate (Optional[float]) – Override starting learning rate for fine-tuning

  • model_branch (str | None)

Returns:

Name of the frozen model file (e.g., ‘frozen_model_2.pth’)

Return type:

str

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:
  • finetune_config (FineTuneConfig) – Fine-tuning configuration from input.yaml

  • datadir (str) – Path to training data directory

  • atom_types (List[str]) – Atomic species list

  • training_dir (str) – Output directory for trained models

  • num_models (int) – Number of ensemble models

  • input_file (str) – DeepMD JSON config (used by DeePMD fine-tuning only)

Returns:

Frozen model filename

Return type:

str

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 created

  • num_models (int) – Number of models to fine-tune (for ensemble/QbC)

  • pretrained_model (str) – Foundation model name (‘small’, ‘medium’, ‘large’) or path to .model file

  • num_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:

str

sparc.src.finetune.setup_MACE_calculator(atoms, model_path)[source]

Setup a MACE calculator for an ASE atoms object.

Parameters:
  • atoms (ase.Atoms) – Atomic structure

  • model_path (str) – Path to the MACE .model file

Returns:

(atoms_with_calc, calculator)

Return type:

tuple