nvalchemi-fine-tuning

nvalchemi Fine Tuning

Overview

Use FineTuningStrategy when adapting pretrained weights to a new dataset, objective, trainable parameter set, or model head. Link users to docs/userguide/finetuning.md, docs/userguide/training.md, docs/userguide/models.md, and docs/userguide/losses.md for full details.

import torch

from nvalchemi.training import (
    CheckpointHook,
    EnergyMSELoss,
    FineTuningStrategy,
    ForceMSELoss,
    OptimizerConfig,
    ValidationConfig,
    create_model_spec,
    default_training_fn,
)

CLI Usage

Use nvalchemi-training finetune when the user wants quick experimentation: an offline JSON spec, a scaffold for a supported source model, a Rich intent report, or direct CLI execution without needing full API knowledge. Use a Python script with FineTuningStrategy when the user needs arbitrary code, custom model construction, dynamic data routing, dynamic losses, or non-standard orchestration. Use nvalchemi-training train init for training-from-scratch specs. The main groups are train, finetune, schema (dump, template), and spec (report, run). Fine-tuning sources live under finetune init: checkpoint, mace, aimnet2, and custom.

Common flow:

nvalchemi-training finetune init mace small-0b \
  --dataset data/train.zarr \
  --output-dir runs/mace-ft \
  --loss-dtype-policy prediction_to_target \
  --out mace-ft.json
nvalchemi-training spec report mace-ft.json
nvalchemi-training spec run mace-ft.json

Use --loss-dtype-policy on finetune init or train init when the CLI scaffold should serialize dtype alignment in strategy.loss_fn_spec. spec report renders the selected policy before execution.

Repeat --dataset to record a MultiDataset workflow. Use torchrun ... -m nvalchemi.training.cli spec run SPEC --distributed for DDP; the CLI initializes DistributedManager, prepends DDPHook, builds the dataset(s), constructs the strategy, and calls run(...).

Runtime hooks belong in source.hooks. Each entry contains a spec object that is the serialized BaseSpec itself: cls_path, timestamp, and the constructor keyword fields for the hook. The CLI builds the hook during spec validation and rejects entries that are not Hook or CheckpointableHook instances. The optional stages list uses TrainingStage names to override where the hook fires, and spec report lists hook firing order chronologically. For model-input transforms such as neighbor lists, use BEFORE_FORWARD; this stage is reused by training and strategy-owned validation. Do not add a validation-only callback for this.

from nvalchemi.hooks import NeighborListHook
from nvalchemi.models.base import NeighborConfig
from nvalchemi.training import create_model_spec

hook_entry = {
    "spec": create_model_spec(
        NeighborListHook,
        config=NeighborConfig(cutoff=5.0),
    ).model_dump(mode="json"),
    "stages": ["BEFORE_FORWARD"],
}

Expect spec report to include warnings for common mistakes such as high fine-tuning learning rates, missing validation data, unsafe checkpoint output paths, or MACE compile settings.

Choose The Entry Point

• Use FineTuningStrategy(models=...) when the user already loaded or built a trainable model.
• Use FineTuningStrategy.from_pretrained_checkpoint(...) to start a fresh fine-tuning run from model weights in a native nvalchemi checkpoint.
• Use FineTuningStrategy.load_checkpoint(...) only to resume an interrupted fine-tuning run with optimizer/scheduler/counters/hook state.

from_pretrained_checkpoint loads the complete checkpoint model set. A single-model checkpoint becomes a single model input; multi-model checkpoints preserve their named mapping. Source optimizer state, hooks, validation settings, counters, and num_epochs/num_steps do not carry over. If the user omits loss_fn or optimizer_configs, they may opt into source metadata with use_original_loss=True or use_original_opt_class=True. Reused optimizer configs get optimizer_lr=1e-5 by default; pass optimizer_lr=None to keep the checkpoint LR.

Minimal Pattern

loss_fn = EnergyMSELoss() + ForceMSELoss(normalize_by_atom_count=True)
loss_fn.dtype_policy = "prediction_to_target"  # optional dtype alignment

strategy = FineTuningStrategy(
    models=pretrained_model,
    trainable_patterns=("main.model.readout.*",),
    optimizer_configs=OptimizerConfig(
        optimizer_cls=torch.optim.AdamW,
        optimizer_kwargs={"lr": 3e-4, "weight_decay": 1e-6},
    ),
    training_fn=default_training_fn,
    loss_fn=loss_fn,
    validation_config=ValidationConfig(validation_data=val_loader, every_n_epochs=1),
    hooks=[CheckpointHook("runs/finetune/checkpoints", epoch_interval=1)],
    num_epochs=10,
)
strategy.run(train_loader)

Use low learning rates for full-model fine-tuning. Prefer trainable_patterns for head-only or adapter-style workflows; patterns match fully qualified names such as "main.model.readout.weight".

From A Native Checkpoint

Use this when a previous nvalchemi run produced a restartable checkpoint but the new task should get fresh fine-tuning counters and optional source loss/optimizer metadata.

strategy = FineTuningStrategy.from_pretrained_checkpoint(
    "runs/pretrain/checkpoints",
    use_original_loss=True,
    use_original_opt_class=True,
    optimizer_lr=1e-5,
    training_fn=default_training_fn,
    trainable_patterns=("main.model.readout.*",),
    num_steps=2_000,
)

For multi-model checkpoints, write training_fn(models, batch) and pass optimizer_configs keyed by the model(s) to update. Models omitted from optimizer_configs are frozen/eval during training but can be used as teachers or references. use_original_loss and use_original_opt_class require native strategy metadata; they do not work with component-only checkpoints.

Bring Your Own Model Or Foreign Checkpoint

Prefer native wrapper constructors for supported pretrained models, for example MACEWrapper.from_checkpoint(..., compile_model=False), because they preserve reconstruction metadata for later strategy checkpoints. compile_model=True is inference-only for MACE and freezes parameters.

For arbitrary PyTorch checkpoints:

• Instantiate the architecture through a wrapper class when possible.
• Use create_model_spec(wrapper_cls_or_factory, ...) for reproducible rebuilds.
• Load weights with state_dict; use strict=False only for intentional head or adapter changes and inspect missing/unexpected keys.
• If output keys differ from loss keys, write a training_fn that returns the mapping expected by the loss.
• Treat foreign checkpoints as weight imports, not restart checkpoints. Save a fresh FineTuningStrategy checkpoint before relying on resume behavior.

state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
model.load_state_dict(state["model"] if "model" in state else state, strict=False)
model_spec = create_model_spec(MyWrapper.from_pretrained, checkpoint_path=str(checkpoint_path))

Patch Or Freeze The Model

Use module_patches to replace or add child modules before optimizer construction. Use create_model_spec(...) for patches that must serialize; direct module instances are runtime-only.

strategy = FineTuningStrategy(
    models=pretrained_model,
    module_patches={
        "main.model.readout": create_model_spec(
            torch.nn.Linear,
            in_features=128,
            out_features=1,
        )
    },
    freeze_patterns=("main.model.*",),
    trainable_patterns=("main.model.readout.*",),
    optimizer_configs=OptimizerConfig(optimizer_cls=torch.optim.AdamW),
    training_fn=default_training_fn,
    loss_fn=EnergyMSELoss(),
    num_steps=1_000,
)

trainable_patterns alone is an allow-list. freeze_patterns excludes broad regions first, then trainable_patterns re-includes exceptions. Use freeze_mode="optimizer_only" only when frozen parameters should still receive gradients for diagnostics or custom hooks.

Typical strategies to fine-tune without catastrophic forgetting include adding different readout/output heads or a new atom embedding table. Users will likely need a way to route based on dataset. If the user does not specify a strategy, discuss options tailored to the model and fine-tuning dataset. Note that equivariant models like MACE will need specialized read-out layers as to preserve equivariance.

Caveats

• Check target/prediction units, atom ordering, neighbor-list assumptions, PBC, dtype, device, and output shapes before training. If label and model-output dtypes differ intentionally, make the user aware of dtype_policy: use "prediction_to_target" to cast outputs to labels or "target_to_prediction" to cast labels to outputs. Set it on a leaf loss, on ComposedLossFunction(...), or as loss_fn.dtype_policy = ... after operator-sugar construction.
• Enable force/stress outputs in the model config when those losses need autograd-derived quantities.
• Start with validation and short checkpoint intervals; pretrained runs can regress quickly with mismatched data or too-large learning rates.
• Resume interrupted fine-tuning with FineTuningStrategy.load_checkpoint, not from_pretrained_checkpoint.