用途
nvalchemi Distillation 用于训练一个小型学生模型，使其复现一个大型冻结教师模型（MLIP，机器学习势）的输出。在 nvalchemi 中，它是 TrainingStrategy 的子类，作用于两个命名模型（"student" 与 "teacher"），因此 nvalchemi-training-api 中的所有能力均适用：优化器、调度器、验证、钩子、检查点、DDP。蒸馏新增的是教师模型，其输出成为损失目标。
主要能力
训练循环
- 离线：在固定数据集上训练，教师标签可预计算到存储中或即时生成。适合起步，像普通训练一样跨 rank 分布。
- On-policy：用学生自身运行动力学生成帧，教师标注，然后在生成帧与参考数据的混合上训练。适用于学生在参考分布上稳定但在实际访问的状态上失败的情况。
教师信号
InProcessTeacherScorer 将 Batch 转换为命名信号，每个信号映射到批次字段：energy → teacher_energy（系统级）、forces → teacher_forces（节点级）、stress → teacher_stress（系统级）、atomic_energies → teacher_atomic_energies（节点级）、embeddings → teacher_node_embeddings（节点级，需额外前向）、hessian → teacher_hvp + teacher_hvp_probe（节点级，需额外能量前向和两次反向）。支持 TeacherSignal 自定义规格，字段名须以 teacher_ 开头。评分过程会收窄教师 active_outputs、构建并回滚教师邻居列表、分离所有输出，保证输入批次原样返回。
离线标注
label_dataset 按连续块遍历数据集，附加教师字段，写入 Zarr 存储，可通过普通 AtomicDataZarrReader / Dataset 路径读取。支持断点续跑（resume=True），丢弃临时邻居张量，拒绝模式漂移的块和不一致的存储。批次未标注时由内部 BEFORE_FORWARD 钩子即时标注。
损失函数
任何内置损失项通过将 target_key 指向教师字段即可蒸馏。AtomicEnergyMatchingLoss 匹配教师逐原子能量分解。ComposedLossFunction 默认重归一化权重，可设 normalize_weights=False 使用字面系数。
高级目标
EmbeddingMatchingLoss 需要 training_fn=embedding_distillation_fn，宽度不同时需注册 EmbeddingProjector 为第三个模型 "projector"。HessianMatchingLoss 需要 training_fn=hessian_distillation_fn 和 hessian 信号，其图平衡值比力 MSE 高一到两个数量级，建议初始比力项轻 100 到 10000 倍。BoltzmannMatchingLoss 要求 on_policy，拒绝 relaxation 传播器和任何收敛准则，拒绝用于验证损失，replay_ratio < 1 时警告。
On-Policy 生成
OnPolicyConfig 描述一个生成-标注-训练段。设置 on_policy= 后 run() 不接受 dataloader。传播器必须持有 models["student"] 注册的同一模块对象。reference_dataset 是混合中参考份额，除非 replay_ratio == 1 否则必需。initial_structures 是任意 StructureSource，InitialStructures 为参考源。OnPolicySettings.probe 默认 True，在构造时对一个初始结构运行传播器 compute() 和 relaxation 准则。divergence 为运行时谓词，默认 nonfinite_divergence。毕业后在 ON_GRADUATE 报告。capture_sink 最多增长一次。
多卡运行
数据并行，添加 DDPHook 每 GPU 启动一个进程。每个 rank 传播自己的初始结构分片，用本 rank 的教师副本标注，填充自己的重放缓冲区；学生梯度 all-reduce 是 rank 间唯一的每步训练通信。初始结构大小应为 world 的整数倍并按原子数排序。多 rank 启动若学生未包装则拒绝。rank 通过 rank_seed_stride（默认 1_000_003）去相关。
检查点与重启
检查点序列化所有模型（含教师），因此周期性 CheckpointHook 写入成本包括教师权重。传播器状态不检查点化；检查点携带计数器、权重、优化器状态和可检查点化的钩子状态。
使用方式
请参考官方文档 docs/userguide/distillation.md（概念指南）和 docs/modules/training/distillation.rst（逐符号参考）。离线最小模式示例：创建 DistillationStrategy 时传入 models={"student": student, "teacher": teacher}、optimizer_configs 仅含 "student" 条目、loss_fn 指向 teacher_* 字段，然后调用 strategy.run(dataloader)。on-policy 模式需额外构造 OnPolicyConfig 并设置 reference_dataset 和 on_policy 参数，然后调用 strategy.run()。
适用场景
适用于将大型教师 MLIP 蒸馏为小型学生模型，尤其是学生需在自身访问的状态上改进的 on-policy 场景，以及需要利用教师嵌入、Hessian 或 Boltzmann 分布匹配等无参考数据列的目标。
限制
- 教师模型通过省略 optimizer_configs 中的 "teacher" 条目冻结；添加教师优化器条目是错误。
- 模型名称固定为 "student" 和 "teacher"。
- 教师信号由损失项推导，除非显式设置 teacher_signals= 请求无损失消耗的信号。
- on-policy 下传播器必须持有学生模块对象本身。
- reference_dataset 必须为教师标注的数据集且为 replay 帧形状，携带自身参考能量或力的数据集会被拒绝。
- 已携带状态迁移器、多子阶段 FusedStage 或 DomainParallel 的传播器在构建 OnPolicyConfig 时被拒绝。
- 检查点成本包含教师权重。
- 传播器状态不检查点化。
- 更多具体限制需参考官方文档。