在大模型训练工程中，「Megatron → HF → 训练 → Megatron」是行业主流闭环：Megatron 负责超大规模预训练，HF 负责微调/对齐，再转回 Megatron 做推理/部署。本文用一条时间线串起数据采样、Tokenization、格式差异、训练参数选取、格式互转等核心环节，并给出可直接落地的脚本与可视化例子。

---

一、为什么要做「数据采样 + 拼接(concat)」

1. 预训练视角：多源异构 + 长度极长

• 网页、代码、论文、书籍等来源比例差异大，需 按 token 比例重采样 以避免低资源语料被淹没
• 单篇文档常 >100 k token，而 GPU 实际能消化的 样本长度（seq_len）一般 2k/4k/8k；必须 拼接 + block 分块 才能填满显存，提高 MFU

2. 微调视角：多样性 + 均衡

• 指令数据往往「长问答少、短问答多」→ 不加采样会出现「短样本垄断 batch，长样本 OOM」
• 分类任务需 类别均衡采样；对话数据需 多轮-单轮混合采样，防止模型「只记得最后一轮」

3. 采样策略代码实战（HF 风格）

from datasets import load_dataset, interleave_datasets

web = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True)
code = load_dataset("bigcode/the-stack-dedup", split="train", streaming=True)
math = load_dataset("kentsui/infimath3plus", split="train", streaming=True)

# 按 85% 12% 3% 混合
mix = interleave_datasets(
        [web, code, math],
        probabilities=[0.85, 0.12, 0.03],
        seed=42,
        stopping_strategy="first_exhausted")

---

二、Tokenizer 是什么？用一个小例子看懂原理

1. 定义

Tokenizer =「文本 → 整数列表」的双向映射表；决定模型「看到的」全部符号

2. 以 GPT-2 BPE 为例

原始句子：
"Hel<NAME>!"

分步输出（huggingface/transformers 调试代码）：

from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("gpt2")
text = "Hel<NAME>!"
print(tok.tokenize(text))
# ['Hel', 'lo', ',', 'ĠVLL', 'M', '!']

print(tok(text).input_ids)
# [15496, 4171, 11, 36223, 476, 0]

关键观察

1. 空格被转义成 Ġ 保证「一词一 token」
2. 罕见词 "VLLM" 被拆成两个高频子词 "VL" + "LM"；这就是 BPE 贪心合并 的核心思想
3. 特殊符号：<|endoftext|>(50256) 用于「文档边界」和「padding」

3. 与 Megatron 的差异点

Megatron 训练脚本里需显式提供：

• vocab-file（json 格式，含合并规则）
• merge-file（bpe 合并对）
  HF 把二者打包进 tokenizer.json（快速版）或 tokenizer_config.json，一键加载

---

三、HF 格式长什么样？为何训练阶段偏爱它

1. 目录结构

llama-7b-hf/
├─ config.json          # 模型超参
├─ pytorch_model.bin    # 权重（或 .safetensors）
├─ tokenizer.json       # 分词表 + 解码器
├─ tokenizer_config.json
├─ special_tokens_map.json
└─ generation_config.json

所有文件纯文本或标准张量格式，可直接 torch.load / json.load；调试友好

2. 与 Trainer / PEFT 生态零摩擦

• 支持 load_dataset、DataCollator、Trainer 一条龙
• 支持 bitsandbytes 量化、peft LoRA/AdaLoRA、DeepSpeed 插件，配置即插即用
• 社区 100 k+ 模型统一 HF Hub，微调完一键 push，便于 CI/CD

---

四、训练参数怎么选？（以 7B 继续预训练为例）

参数族	推荐值	说明
seq_len	4096	显存与长文本权衡；A100-80G 可塞 4k×1024 token batch
global_batch_size	2M token	约 0.5B 参数×2k step ≈ 1T token；用 batch_size = 2M / seq_len
lr	3e-4 → cosine ↓ 3e-5	继续预训练用原始预训练学习率×0.3更安全
warmup	2% steps	防止初始冲击
weight_decay	0.1	过滤大参数
grad_clip	1.0	稳定
data_sampling	按 85% 通用 12% 代码 3% 数学 动态混合	见第一节
save&eval	每 0.1 epoch	用 load_best_model_at_end=True 自动回滚

Megatron 侧额外开关

• --tensor-parallel 2 --pipeline-parallel 2
• --micro-batch-size 2 --global-batch-size 1024
• --no-query-key-layer-scaling（与 HF 权重命名对齐）

---

五、训完为何还要转回 Megatron 格式？

1. 推理阶段需要 TP/PP：HF 只支持单卡/单节点，大模型线上必需张量+流水并行
2. Kernel 优化：Megatron 的 fused_layer_norm、fused_attention 比 HF 快 8–15%
3. 量化/剪枝工具链：NVIDIA FasterTransformer、TensorRT-LLM 直接吃 Megatron checkpoint
4. 生产环境统一：多机多卡推理集群（BLOOM-176B、GLM-130B）全部用 Megatron 框架部署

---

六、Megatron vs HuggingFace 格式逐项对比

维度	Megatron-LM	HuggingFace
权重切片	按 TP/PP 拆成 mp_rank_00/model_optim_rng.pt	整图 pytorch_model.bin
键名	language_model.encoder.layers.*	model.layers.*
分词表	独立 vocab.json + merge.txt	内置 tokenizer.json
超参	args.json + 脚本入参	config.json 统一
并行策略	显式 TP/PP/DP	用 device_map 隐式
继续训练	需重走 preprocess_data.py → 索引	load_dataset 直接流式
生态插件	DeepSpeed-MII、FasterTransformer	PEFT、TRL、Optimum

---

七、互转实战：脚本 + 可视化流程

1. HF → Megatron（Pre-train / TP=2）

# Step1 安装官方转换器
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
pip install -r requirements.txt

# Step2 下载 HF 模型
huggingface-cli download meta-llama/Llama-2-7b-hf --local-dir ./llama-7b-hf

# Step3 转换 + 自动切片
python tools/checkpoint_util.py \
        --model-type GPT \
        --load-dir ./llama-7b-hf \
        --save-dir ./llama-7b-megatron-tp2 \
        --tensor-model-parallel-size 2 \
        --target-tensor-model-parallel-size 2 \
        --hf-to-megatron

输出目录：

llama-7b-megatron-tp2/
├─ mp_rank_00/model_rng.pt   # GPU0 切片
├─ mp_rank_01/model_rng.pt   # GPU1 切片
└─ common_args.json          # 新超参

2. Megatron → HF（微调后合并）

# 假设 TP=2 PP=1 微调结束
python tools/checkpoint_util.py \
        --model-type GPT \
        --load-dir ./llama-7b-megatron-tp2 \
        --save-dir ./llama-7b-hf-output \
        --tensor-model-parallel-size 2 \
        --target-tensor-model-parallel-size 1 \
        --megatron-to-hf

原理：把 mp_rank_00 与 mp_rank_01 的列并行权重沿隐藏维度 cat 回来，再重命名键到 HF 规范

3. 一键双向转换函数（Python API）

from megatron.checkpointing import load_checkpoint, save_checkpoint
from transformers import LlamaForCausalLM, LlamaConfig
import torch

def megatron_to_hf(mega_dir, hf_dir, tp_size=2):
    """把 Megatron TP 切片 → HF 整图"""
    config = LlamaConfig.from_pretrained(hf_dir)          # 读目标超参
    model = LlamaForCausalLM(config)
    # 1. 加载 TP 切片
    all_weights = {}
    for tp in range(tp_size):
        ckpt = torch.load(f"{mega_dir}/mp_rank_{tp:02d}/model_rng.pt", map_location="cpu")
        all_weights[tp] = ckpt["model"]["language_model"]["encoder"]
    # 2. 合并列并行（MLP.down_proj / Attention.o_proj）
    merged = merge_tensor_parallel(all_weights, tp_size)
    # 3. 键映射 + 写回
    model.load_state_dict(merged, strict=False)
    model.save_pretrained(hf_dir)

def hf_to_megatron(hf_dir, mega_dir, tp_size=2):
    """HF 整图 → Megatron TP 切片"""
    model = LlamaForCausalLM.from_pretrained(hf_dir, torch_dtype=torch.float16)
    # 1. 切分列并行
    sliced = split_tensor_parallel(model.state_dict(), tp_size)
    # 2. 逐个保存
    for tp, shard in enumerate(sliced):
        save_checkpoint(f"{mega_dir}/mp_rank_{tp:02d}", shard)

def merge_tensor_parallel(weights, tp):
    """示例：合并 MLP.down_proj 列并行"""
    keys = ["mlp.down_proj.weight"]          # 行并行直接 concat
    merged = {}
    for k in keys:
        merged[k] = torch.cat([weights[i][k] for i in range(tp)], dim=0)
    return merged

---

八、小结 & 一张图记住全流程

原始语料
   ├─ 采样&拼接 → 均衡混合 + 最长上下文填充
   ├─ Tokenizer → 文本变 token ID（BPE/SentencePiece）
   ├─ HF格式 → 训练/微调（PEFT + Trainer）
   └─ Megatron格式 → TP/PP 推理加速
         ↑  互转脚本（checkpoint_util / 自定义API）
         └─ 生产部署（FasterTransformer / TensorRT-LLM）

记住三句话

1. 数据先「采样」再「拼接」→ 显存不浪费，学习曲线平滑
2. HF 格式是「训练期通用语言」；Megatron 格式是「推理期并行语言」
3. 转换本质是「键名重映射 + 张量合并/切片」，官方脚本已覆盖 90% 场景