工具框架(二)：NVIDIA NeMo / TensorRT Model Optimizer

工具篇第二篇聚焦工业级大模型蒸馏链路：NVIDIA NeMo Framework 与 TensorRT Model Optimizer（ModelOpt）。前置：你应理解 KD 基本概念（Hinton et al. 2015），并最好具备大模型（GPT 类）训练与并行基础。本文依据 NVIDIA 官方 NeMo Framework 蒸馏文档（docs.nvidia.com/nemo-framework）与 ModelOpt 能力说明展开。

一、NeMo 中的蒸馏能力

NeMo 2.0 提供了”开箱即用”的知识蒸馏训练设置，其蒸馏能力由 TensorRT Model Optimizer（ModelOpt）启用——ModelOpt 是 NVIDIA 用于”在 GPU 上优化深度学习模型以服务推理”的库。官方文档明确：ModelOpt 支持训练时蒸馏与量化感知训练（见 NeMo Framework 官方蒸馏文档），这意味着蒸馏可以和质量压缩（量化）放在同一训练流程里。

蒸馏的两个核心价值（官方表述）：相比从零传统训练，收敛更快、最终精度更高。

二、Logits 蒸馏流程

NeMo 的 logits 蒸馏过程可概括为四步：

加载检查点：同时加载学生与教师检查点；两者必须支持相同的并行策略（如张量并行 TP 大小一致）；
替换损失：把标准损失替换为输出 logits 之间的 KL 散度（也可加入若干中间层状态之间的附加损失）；
训练：对学生与教师都跑前向，但反向传播只作用于学生；
保存：仅保存学生检查点，后续可像普通模型一样使用。

伪代码式的训练步：

for batch in dataloader:
    t_logits = teacher(**batch).logits          # no_grad
    s_logits = student(**batch).logits
    loss = KL_div(softmax(t_logits/T), softmax(s_logits/T)) * (T*T)
    loss.backward()                             # 仅学生参数更新
    optimizer.step()

这与 Hinton et al. 2015 的软损失一致，只是工程上扩展到大模型并行训练。

三、配置与中间层蒸馏

NeMo 通过 YAML 配置 distillation 行为，典型字段：

logit_layers: ["output_layer", "output_layer"]
intermediate_layer_pairs:
  - ["decoder.layers.3",  "decoder.layers.3"]
  - ["decoder.layers.6",  "decoder.layers.9"]
  - ["decoder.layers.11", "decoder.layers.18"]
skip_lm_loss: true
kd_loss_scale: 1.0

logit_layers：学生/教师输出 logits 层名；
intermediate_layer_pairs：成对的中间层，默认用余弦相似度损失对齐（可视为特征级蒸馏，承袭 FitNets / Romero et al. 2015 的 hint 思想）；
skip_lm_loss：是否跳过原始语言建模损失；为 false 时 LM 损失会加到蒸馏损失上；
kd_loss_scale：蒸馏损失相对 LM 损失的缩放系数。

四、与量化感知、剪枝协同

ModelOpt 的”训练时蒸馏 + 量化感知”能力，使你可以在一次训练中同时逼近低精度与教师分布。对于大模型，常见范式是先结构化剪枝、再用蒸馏恢复精度（NVIDIA Minitron 式实践）：剪掉若干层/宽度后，以原模型为教师做蒸馏微调。关于该范式节省的具体训练 token 倍数，NVIDIA 文档有公开数字，但不在本事实库范围内，本文标记为（待核实），请以官方文档为准。

五、使用方式与局限

NeMo 提供两种入口：

NeMo-Run Recipe：from nemo.collections.llm.modelopt.recipes import distillation_recipe，填入学生/教师检查点路径与配置即可发起；
torchrun / Slurm 脚本：直接调用 scripts/llm/gpt_train.py，用 model_path + teacher_path 指定学生与教师，适合多节点大规模训练。

需注意的局限（以官方文档为准，随版本演进）：当前主要支持 GPT 类 NeMo 2.0 检查点；部分版本仅启用 logit-pair 蒸馏；启用流水线并行（PP）时，中间层蒸馏损失仅支持在最后一个 pipeline stage。HF 模型可经检查点转换器先转 NeMo，蒸馏后再转回 HF 格式。