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(...).