natural_language_autoencoders（NLA，Natural Language Autoencoders）是 Anthropic 可解释性研究「用自然语言解释 LLM 激活向量」的完整训练仓库。它把大模型残差流向量变成自然语言描述（AV），再从描述还原向量（AR）。训练完成后，你面对的通常是一堆FSDP / Megatron 分布式检查点，而推理（如 SGLang 服务）需要Hugging Face 格式，部署大模型还常要FP8 / INT4 / MXFP8 量化。NLA 在 tools/ 目录下提供了 8 个开箱即用的检查点转换脚本，覆盖「训练格式 ↔ HF 格式」互转与 3 种量化的完整闭环。本文带你逐一掌握这套工具链。

一、为什么需要检查点转换？

NLA 的训练栈构建在 Miles（FSDP2 / Megatron 后端）之上：

• 7B / 12B / 27B 模型用 FSDP 后端训练，产出的检查点是 DCP（torch.distributed.checkpoint）格式的iter_xxx目录，HF 推理框架无法直接加载；
• Llama-70B 用 Megatron 后端（TP + PP），产出的检查点是torch_dist格式，按层/专家/并行维度切分；
• 推理与发布需要标准 HFsafetensors目录（config.json+model-xxxx.safetensors+ tokenizer），每个检查点还附带nla_meta.yamlsidecar 记录 prompt 模板、注入 token ID 与缩放因子。

因此「检查点转换」是 NLA 工作流中训练和推理之间的必经环节。

二、tools/ 工具链总览：8 个脚本一张表

三、FSDP → HF：convert_fsdp_to_hf.py（最常用路径）

Miles FSDP 训练（AV 的 SFT / GRPO RL）保存的iter_xxx目录就是 DCP 检查点。转换命令只需三个参数：

python tools/convert_fsdp_to_hf.py \ --input-dir /path/to/actor/iter_0004000 \ --output-dir /path/to/actor_hf/iter_0004000 \ --origin-hf-dir Qwen/Qwen2.5-7B-Instruct \ --force

参数说明（见 tools/convert_fsdp_to_hf.py）：

• --input-dir：DCP 检查点目录（可传iter_xxx或iter_xxx/model，会自动探测）；
• --output-dir：输出 HF 目录，已存在时需加--force覆盖；
• --origin-hf-dir：原始 HF 模型目录，用于提供 config / tokenizer 等资产；
• 脚本还会把 tokenizer 等非权重文件自动拷贝到输出目录（copy_assets）。

三个值得了解的细节：

1. 键名前缀自动剥离：DCP 键名常带model_state./module.等包装前缀，工具会枚举候选前缀并选取与 HF 骨架匹配数最多者（_strip_best_prefix）。
2. Gemma 多模态包装器自动识别：Gemma-3 的 config 带text_config，而 NLA 实际保存的是纯文本 CausalLM，工具会自动改用文本子配置构建骨架，避免「静默加载成随机初始化」的坑（_resolve_skeleton_config）。
3. dtype 修复：DCP 本身是 fp32，骨架模型的torch_dtype可能为空，SGLang 加载会报 dtype 不匹配。工具会从 DCP 张量取真实 dtype 并写回 config（tools/convert_fsdp_to_hf.py）。

⚠️ 注意：转换产物可能是 fp32 体积（约 2 倍），下游加载时可按需torch_dtype=torch.bfloat16，NLA 的 ref model 加载就是如此处理的（nla/train_actor.py）。

四、Megatron ↔ HF：Llama-70B 级别的三条转换路径

Megatron 路线用于 Llama-70B 级别的训练，涉及「HF 初始化 → torch_dist 保存 → 转回 HF」三个方向。

1. HF → Megatron torch_dist：convert_hf_to_torch_dist.py

训练前的准备脚本：从 HF 基座模型构建 Megatron 分布式检查点。

• 通过mbridge把 HF 权重加载进 Megatron 模型并保存为torch_dist格式；
• 自动按层数推算 pipeline 并行度（world_size必须 ≤ 层数）；
• 打了一个FP32 保精度补丁：_weight_to_mcore_format在转换_keep_fp32参数时保持原精度（tools/convert_hf_to_torch_dist.py）；
• 支持 ROCm（自动替换FileSystemWriterAsync）。

NLA 专属参数--nla-model-is-critic（源码）：AR（critic）模型把 Megatron 的output_layer替换为[d, d] + Identity，使torch_dist检查点与训练时的 model_provider 结构一致。配套的截断检查点由 nla/scripts/prepare_critic_checkpoint.py 生成——它把基座模型截断到第 K 层（去掉lm_head与最终 LN，value_head用单位阵初始化），并可加--megatron-compat写入 mbridge 兼容键（--megatron-compat 说明）。

2. Megatron → HF（单机离线）：convert_torch_dist_to_hf.py

训练结束后把torch_dist目录离线转成 HFsafetensors，无需多卡：

python tools/convert_torch_dist_to_hf.py \ --input-dir /path/to/llama70b/iter_0024000 \ --output-dir /path/to/llama70b_hf/iter_0024000 \ --origin-hf-dir Llama/Llama-3.3-70B-Instruct \ --vocab-size 128256 \ --force

要点（tools/convert_torch_dist_to_hf.py）：

• 从检查点的common.pt读取 Megatron 训练参数（层数、专家数等），把合并存储的层参数 / MoE 专家参数逐个拆开还原；
• --vocab-size用于去掉 Megatron 的词表 padding 行；
• 输出按--chunk-size（默认 5GB）切分为多个model-xxxxx-of-yyyyy.safetensors，并生成model.safetensors.index.json。

3. Megatron → HF（多卡在线）：convert_to_hf.py

权重留在 GPU 上按TP / PP / EP 并行组在线 gather后转换，适合超大模型单机放不下磁盘中间态的场景。基于 Miles 的 Megatron 后端运行，新增参数：

• --output-dir：保存转换后的 HF 模型；
• --check-same：与原始 HF 模型逐参数比形状和数值，校验转换正确性（--check-same 逻辑）——建议转换后先跑一次校验再上线。

此外，NLA 训练中的 critic（AR）模型有专门的在线 HF 导出路径 nla/megatron/checkpoint.py：复用 SGLang 权重同步的 gather 管线，把lm_head.weight重定向为value_head.safetensors，产出与NLACriticModel.from_pretrained完全兼容的 HF 目录。

五、模型瘦身：FP8 / MXFP8 / INT4 量化转换

转成 HF 格式后，若要部署 70B 级模型，可继续在 tools/ 里做量化。

1. BF16 → FP8：convert_hf_to_fp8.py

python tools/convert_hf_to_fp8.py \ --model-dir /path/to/hf_bf16 \ --save-dir /path/to/hf_fp8 \ --strategy block \ --block-size 128 128 \ --max-workers 4

• --strategy三种量化粒度（tools/convert_hf_to_fp8.py）：
  • block（默认）：分块量化，--block-size指定块大小，精度与压缩率最均衡；
  • channel/tensor：逐通道 / 整张量一个 scale，输出走 compressed-tensors 格式；
• 自动跳过不宜量化的模块（layernorm、embedding、router、mlp.gate、lm_head 等），并写入modules_to_not_convert；
• 转换结果写回config.json的quantization_config并重建model.safetensors.index.json，SGLang / vLLM 可直接加载；
• 多线程处理分片（--max-workers）。

2. BF16 → MXFP8：convert_hf_to_mxfp8.py

python tools/convert_hf_to_mxfp8.py \ --model-dir /path/to/hf_bf16 \ --save-dir /path/to/hf_mxfp8

• MXFP8 采用1×32 分组 + UE8M0 缩放因子（quantization_config），scale 布局与 SGLang 的 Triton 量化内核镜像对齐，可直接被推理引擎消费；
• 依赖 SGLang 的mxfp8_group_quantize，且必须在 CUDA 设备上运行；
• 自动跳过 1D 权重、末维不被 32 整除的张量与 norm / embed 等敏感模块。

3. BF16 → INT4：convert_hf_to_hf_int4.py

基于 LLM Compressor 的GPTQ离线量化，支持--quant-type W4A16 / W8A16：

• 需要提供--data-dir校准数据集（读取其中train-00000-of-00001.parquet的text列，随机抽--num-calibration-samples条、每条--max-sequence-length个 token）；
• --quant-group-size默认 32、--dampening-frac默认 0.01（GPTQModifier 配置）；
• 自动忽略lm_head、norm、embedding、注意力投影等模块。

💡 校准数据建议用与目标任务同分布的文本；NLA 训练本身就产出 parquet 语料，可复用其 Stage 0 提取前的原始文本。

4. FP8 → BF16：fp8_cast_bf16.py（反向去量化）

目标环境不支持 FP8 内核时，把 FP8 检查点还原为 BF16：

python tools/fp8_cast_bf16.py \ --input-fp8-hf-path /path/to/hf_fp8 \ --output-bf16-hf-path /path/to/hf_bf16

用 Triton 内核按_scale_inv反量化 FP8 权重（weight_dequant_kernel），同步更新 index 文件并复制 config / tokenizer。

六、实操技巧与常见坑位 📌

1. 先校验再发布：Megatron 路径转换后用--check-same比对原始 HF 模型；FSDP 路径转换后重点看日志里的Missing keys（非空即说明骨架不匹配，权重是随机初始化）。
2. fp32 体积翻倍：DCP 原生 fp32，转出的 HF 检查点可能比 BF16 大一倍；下游推理侧显式指定torch_dtype=torch.bfloat16加载即可。
3. Gemma-3 用户：FSDP 转换脚本已自动处理多模态包装器问题，无需手工改 config；但推理时记得 Gemma 的 embed 缩放与--attention-backend fa3等特殊处理（详见 docs/inference.md）。
4. critic（AR）模型不是标准 CausalLM：它没有lm_head/ 最终 LN，value_head独立存放。初始化 Megatron 检查点时务必加--nla-model-is-critic，配合--megatron-compat生成的兼容键。
5. 保留 sidecar：发布检查点必须带nla_meta.yaml（prompt 模板、注入 token ID、缩放因子），公开前用 release/hf_stage_and_scrub.py 清洗内部路径。
6. 输出目录冲突：所有转换工具都要求--output-dir / --save-dir不存在（FSDP / torch_dist 路径用--force覆盖），脚本失败重试时先清理残留目录。

七、延伸阅读

• docs/design.md：NLA 训练系统设计（FSDP / Megatron 双后端、AV/AR 双模型调度）
• docs/setup.md：Miles + SGLang + NLA 环境安装
• docs/inference.md：转换后检查点的 SGLang 推理与 AR 打分
• configs/TRAINING_NOTES.md：各规模模型训练配置与超参笔记
• examples/：四个已发布检查点的真实解码示例
• release/hf_stage_and_scrub.py：检查点发布前的暂存与元数据清洗

掌握以上 8 个脚本后，你就能在「Miles 训练产出 → HF 推理部署 → 低精度量化上线」之间自由切换，完整跑通 NLA 从训练到发布的全链路。