Skip to content

Latest commit

 

History

History
195 lines (147 loc) · 7.27 KB

File metadata and controls

195 lines (147 loc) · 7.27 KB
name deepmd-finetune-dpa4
description Fine-tune a DPA4 model in DeePMD-kit. Use for standard or LoRA fine-tuning from a DPA4/SeZM .pt checkpoint, validation and .pt2 export.
compatibility Requires deepmd-kit with the PyTorch backend. DPA4/SeZM training is GPU-oriented.
license LGPL-3.0-or-later
metadata
author version repository
SchrodingersCattt
1.0

DeePMD-kit Fine-tuning: DPA4

Fine-tune a DPA4/SeZM checkpoint on downstream DeePMD data. This skill covers single-task standard and LoRA fine-tuning. Do not infer the model family from a .pt suffix or filename: DPA3 and DPA4 checkpoints use the same suffix.

Route the checkpoint

If the user has not already established the model family, inspect the stored configuration:

dp --pt show pretrained.pt descriptor fitting-net type-map

Use this skill only when the descriptor/model configuration identifies DPA4 or SeZM. If the checkpoint is multi-task, inspect its branches before selecting a head:

dp --pt show pretrained.pt model-branch descriptor type-map

Do not guess a branch. Use deepmd-finetune-dpa3 instead when the descriptor is DPA3, and stop when the family cannot be established.

Obtain a pretrained checkpoint

Fine-tuning requires a DPA4/SeZM training checkpoint (.pt), not a .pt2 deployment archive. Check whether the installed version provides one:

dp pretrained download -h

Use only a listed model or a checkpoint supplied by the user or its publisher. Record its source and DeePMD-kit version, then verify its descriptor, branch, architecture, and type_map before use.

Before fine-tuning

  1. Confirm the checkpoint exists and can be inspected.
  2. Confirm training, validation, and held-out systems, labels, and element type maps.
  3. Split correlated frames by independent system, trajectory, or source family; do not create a nominal held-out set by randomly splitting adjacent frames.
  4. Validate each DeePMD system before training: natoms is the number of tokens in type.raw, coordinate and force widths are 3 * natoms, and every used label is finite and frame-aligned.
  5. Start from the exact checkpoint architecture. Introducing new element types, changing architecture, or combining specialized spin/property/multi-task configurations requires separate compatibility validation.
  6. Choose standard fine-tuning or LoRA. Do not assume a built-in DPA4 model name; check dp pretrained download -h for the installed version.

Decide whether to use LoRA

Use standard fine-tuning by default. Use LoRA only for a single-task target when parameter-efficient adaptation is wanted and the exact base architecture is known. LoRA is enabled by a non-null model.lora block in the new input; the pretrained checkpoint does not need to contain LoRA. Multi-task LoRA targets are unsupported.

Periodic LoRA checkpoints retain adapters and can resume training. Best checkpoints may merge the adapters into ordinary DPA4 weights, so absence of LoRA metadata does not prove LoRA was never used.

Standard fine-tuning

The model section in input.json must match the checkpoint unless the standard pretrained-script mechanism is deliberately used:

dp --pt train input.json --finetune pretrained.pt

When fine-tuning a single-task target from a multi-task checkpoint and the intent is to preserve a particular pretrained fitting head, pass the branch selected above:

dp --pt train input.json --finetune pretrained.pt --model-branch SELECTED_BRANCH

If --model-branch is omitted, the fitting net may be initialized from the RANDOM branch instead. A multi-task target uses finetune_head in each target branch rather than the command-line option.

If the architecture is unknown, --use-pretrain-script can inherit the stored model configuration except for type_map:

dp --pt train input.json --finetune pretrained.pt --use-pretrain-script

Inspect the resulting configuration and run a bounded initial segment before a long training job. Do not combine model-specific additions with --use-pretrain-script unless that combination has been validated.

LoRA fine-tuning

DPA4/SeZM supports LoRA adapters for single-task fine-tuning. Copy the exact base architecture into lora_ft.json, then add:

{
  "model": {
    "type": "dpa4",
    "lora": {
      "rank": 16,
      "alpha": 16.0
    }
  }
}

Run:

dp --pt train lora_ft.json --finetune pretrained.pt

When pretrained.pt is multi-task, preserve the selected fitting head:

dp --pt train lora_ft.json --finetune pretrained.pt \
    --model-branch SELECTED_BRANCH

Use the shorter command only for a single-task source checkpoint.

The JSON fragment above is not a complete training input. Adapt the full public example at ../../examples/water/dpa4/lora_ft.json, but copy the exact architecture from the source checkpoint before adding model.lora. Do not add --use-pretrain-script unless a targeted test confirms that LoRA is retained.

Monitor and validate

Monitor lcurve.out for non-finite values and train/validation divergence. Select a checkpoint using validation data, then follow the complete held-out evaluation with that exact native checkpoint and every held-out system. Export for deployment only after the complete evaluation meets the task's declared thresholds.

Export and test

DPA4/SeZM uses the .pt2 AOTInductor export path rather than the conventional PyTorch .pth freeze path:

dp --pt freeze -c ckpt/model.ckpt.pt -o finetuned_model
dp test -m finetuned_model.pt2 -s /path/to/test_system -n 30

The freeze command detects DPA4/SeZM and writes finetuned_model.pt2. Validate the exported archive in the target environment before deployment.

For a multi-task checkpoint, freeze the selected head explicitly:

dp --pt freeze -c ckpt/model.ckpt.pt -o finetuned_model --head SELECTED_BRANCH

The resulting .pt2 contains the selected single head; do not pass a branch again when loading that archive.

Checklist

  • The stored descriptor identifies DPA4/SeZM; the .pt suffix was not used as proof.
  • The checkpoint source, DeePMD-kit revision, architecture, and type map are recorded.
  • The intended branch is explicit for a multi-task checkpoint.
  • Training, validation, and held-out systems are independent by source family.
  • Every admitted system has consistent atom counts, shapes, labels, and type maps.
  • The input architecture is compatible with the checkpoint.
  • Standard fine-tuning versus LoRA was selected from the task layout and domain shift.
  • A resumable LoRA checkpoint is distinguished from a merged best checkpoint.
  • LoRA uses a complete base configuration and is not silently overwritten.
  • Complete held-out metrics, sample counts, and reference-label scales are reported.
  • The selected .pt checkpoint was exported to and tested as .pt2.

References