nvalchemi-distillation

nvalchemi Distillation

Overview

Distillation trains a small student to reproduce a large frozen teacher. In nvalchemi it is a TrainingStrategy subclass over two named models, so everything from nvalchemi-training-api applies: optimizers, schedulers, validation, hooks, checkpoints, DDP. What distillation adds is a teacher whose outputs become loss targets.

Read nvalchemi-training-api and nvalchemi-loss-api first. docs/userguide/distillation.md is the concept guide and docs/modules/training/distillation.rst is the symbol-by-symbol reference.

import torch

from nvalchemi.dynamics.integrators.nvt_langevin import NVTLangevin
from nvalchemi.training import (
    ComposedLossFunction,
    EnergyMSELoss,
    ForceMSELoss,
    OptimizerConfig,
)
from nvalchemi.training.distillation import (
    DistillationStrategy,
    InProcessTeacherScorer,
    OnPolicyConfig,
    OnPolicySettings,
    AtomicEnergyMatchingLoss,
    ReplayBuffer,
    InitialStructures,
    TeacherLabelHook,
    build_mixed_loader,
    default_distillation_fn,
    label_dataset,
)

Two loops are available:

• Offline — train over a fixed dataset, with the teacher's labels either precomputed into a store or produced on the fly. Start here; it distributes over ranks like ordinary training.
• On-policy — generate frames by running dynamics with the student itself, label them with the teacher, and train on a mixture of those and reference data. Use it when the student is stable on the reference distribution but fails on the states it actually visits.

Minimal Pattern (offline)

strategy = DistillationStrategy(
    models={"student": student, "teacher": teacher},
    optimizer_configs={
        "student": [
            OptimizerConfig(
                optimizer_cls=torch.optim.AdamW,
                optimizer_kwargs={"lr": 1e-4, "weight_decay": 1e-6},
            )
        ]
    },
    loss_fn=EnergyMSELoss(target_key="teacher_energy")
    + 10.0 * ForceMSELoss(target_key="teacher_forces"),
    num_steps=10_000,
)
strategy.run(dataloader)