test: add GPT2 Megatron baseline

[JYMiracle305]

概述

本 PR 在现有 Megatron shadow baseline 中接入 GPT-2。InfiniTrain、PyTorch 和 Megatron-LM 使用相同 LLMC FP32 初始权重、相同 Tiny Shakespeare token 顺序和相同训练参数，用于比较逐步 loss 与标准训练过程中的吞吐。

GPT-2 baseline 沿用 Llama3 的目录组织、BF16 autocast 语义、10 步基础用例和批量日志比较方式，不修改 third_party/Megatron-LM。运行生成的 dataset、cache 和逐用例日志位于 baseline/megatron/artifacts/；汇总训练日志归档在 baseline/megatron/logs/qy_a100_g3025/。

Stack 关系

master
└── #14 Llama3 Megatron baseline
    └── #15 Qwen3 Megatron baseline
        └── #17 GPT2 Megatron baseline

本 PR 的 base 是 feat/add-qwen3-megatron-baseline。

目录结构

baseline/megatron/
├── logs/qy_a100_g3025/
│   └── result_megatron_training_gpt2.log  # 16 组用例汇总日志
├── models/
│   └── gpt2/
│       ├── README.md                 # 配置、运行与验证说明
│       ├── prepare/
│       │   └── prepare_dataset.sh    # Tiny Shakespeare 数据准备入口
│       ├── train/
│       │   ├── pretrain.py           # LLMC 加载、训练入口、autocast 与性能统计
│       │   ├── run_training.sh       # 单个配置训练入口
│       │   └── run_basic_cases.sh    # 16 个基础用例批量入口
│       └── compare/
│           └── compare_loss.sh       # 单用例 loss/JSON 诊断入口
└── scripts/
    ├── compare_loss.py               # 批量 loss PASS/FAIL 汇总和用例过滤
    └── compare_tps.py                # 批量 TPS 对比

数据与模型对齐

同一份 GPT-2 LLMC FP32 checkpoint
├── InfiniTrain：原生 loader 读取
└── Megatron：校验 LLMC header，恢复 runtime vocabulary，并映射参数布局

同一份 Tiny Shakespeare token stream
├── InfiniTrain：按原始顺序读取 token bin
└── Megatron：使用 GPTDataset，并替换 shuffle-index builder 保持相同顺序

训练结果
├── 单例诊断：models/gpt2/compare/compare_loss.sh
├── 批量精度：scripts/compare_loss.py
└── 批量性能：scripts/compare_tps.py

GPT-2 参数适配会校验 LLMC 文件头、恢复运行时词表大小，并将 QKV 按 attention head 交错为 Megatron MCore 所需布局。LLMC 直接加载覆盖当前基础矩阵中的 DP、TP、SP、PP、VPP 和 3D 并行组合。

BF16 基础用例使用 DTYPE=autocast_bfloat16：参数保持 FP32，forward 位于 PyTorch BF16 autocast 上下文中，与 InfiniTrain/PyTorch 的精度策略对齐。Megatron 原生 DTYPE=bfloat16 仍可用于额外诊断，但不属于默认基础矩阵。FP32 默认关闭 TF32。

基础用例

Case	精度	GPU	DP	TP	PP	SP	VPP	Micro batch	Tokens/step	Sequence	Steps
gpt2_1	FP32	1	1	1	1	关闭	1	4	256	64	10
gpt2_1_bfloat16	BF16 autocast	1	1	1	1	关闭	1	4	256	64	10
gpt2_2	FP32	1	1	1	1	关闭	1	80	5120	64	10
gpt2_2_bfloat16	BF16 autocast	1	1	1	1	关闭	1	80	5120	64	10
gpt2_3	FP32	8	8	1	1	关闭	1	10	5120	64	10
gpt2_3_bfloat16	BF16 autocast	8	8	1	1	关闭	1	10	5120	64	10
gpt2_4	FP32	8	2	4	1	关闭	1	40	5120	64	10
gpt2_4_bfloat16	BF16 autocast	8	2	4	1	关闭	1	40	5120	64	10
gpt2_5	FP32	8	2	4	1	开启	1	40	5120	64	10
gpt2_5_bfloat16	BF16 autocast + TE	8	2	4	1	开启	1	40	5120	64	10
gpt2_6	FP32	8	1	1	8	关闭	1	10	5120	64	10
gpt2_6_bfloat16	BF16 autocast	8	1	1	8	关闭	1	10	5120	64	10
gpt2_7	FP32	4	1	1	4	关闭	2	10	5120	64	10
gpt2_7_bfloat16	BF16 autocast	4	1	1	4	关闭	2	10	5120	64	10
gpt2_8	FP32	8	2	2	2	开启	2	20	5120	64	10
gpt2_8_bfloat16	BF16 autocast + TE	8	2	2	2	开启	2	20	5120	64	10

精度与性能统计

• GPT-2 FP32 loss 最大绝对误差阈值：2e-3
• BF16 loss 最大绝对误差阈值：1e-2
• Megatron iteration 2 的 running-average 日志在比较时恢复为单步 loss
• 共训练 10 步，step 1 作为 warmup，平均 step 2-10
• 多卡使用最慢 rank；TPS 对比报告双方平均 tok/s 和 InfiniTrain/Megatron 百分比

验证方式

bash baseline/megatron/models/gpt2/prepare/prepare_dataset.sh
bash baseline/megatron/models/gpt2/train/run_basic_cases.sh

只运行选定用例：

CASES=gpt2_1,gpt2_8_bfloat16 bash baseline/megatron/models/gpt2/train/run_basic_cases.sh

批量比较全部 16 个用例：

python3 baseline/megatron/scripts/compare_loss.py /path/to/infinitrain/logs/basic baseline/megatron/artifacts/gpt2/logs --include-prefix gpt2_ --threshold-fp32 2e-3
python3 baseline/megatron/scripts/compare_tps.py /path/to/infinitrain/logs/basic baseline/megatron/artifacts/gpt2/logs --include-prefix gpt2_

比较单个用例并输出逐步 JSON：

INFINITRAIN_LOG=/path/to/infinitrain.log MEGATRON_LOG=/path/to/megatron.log OUTPUT_JSON=/path/to/result.json ATOL=1e-2 bash baseline/megatron/models/gpt2/compare/compare_loss.sh

loss 脚本明确列出 PASS、FAIL、缺失日志和总计；TPS 脚本输出平均吞吐、比例、解析错误和缺失日志。任一解析错误或缺失日志都会返回非零退出码。

256-token BF16 小 batch 对运行版本较敏感，可能超过默认 1e-2，仍作为显式失败报告，不静默放宽阈值。

检查

• 16 个 GPT-2 Megatron 基础训练用例完成
• Shell 脚本通过 bash -n
• Python 文件通过编译检查
• git diff --check 通过

[JYMiracle305]

精度对比

rebase 224后，通过用例比之前多了2个

[kilinchange]

当前截图中的 failed loss cases 可能与 InfiniTensor/InfiniTrain#224 修复的问题有关，建议基于 InfiniTrain 最新 master 重新对比后，更新精度对比情况截图。