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/，不会提交到仓库。

Stack 关系

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

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

目录结构

baseline/megatron/
├── models/
│   └── gpt2/
│       ├── README.md                 # 配置、运行与验证说明
│       ├── prepare/
│       │   └── prepare_dataset.sh    # Tiny Shakespeare 数据准备入口
│       ├── train/
│       │   ├── pretrain.py           # LLMC 加载、Megatron 训练入口、autocast 与性能统计
│       │   ├── run_training.sh       # 单个配置训练入口
│       │   └── run_basic_cases.sh    # 六个基础用例批量入口
│       └── compare/
│           └── compare_loss.sh       # 单用例 loss/JSON 诊断入口
└── scripts/
    ├── compare_loss.py               # 批量 loss PASS/FAIL 汇总，增加用例过滤
    └── compare_tps.py                # 继承 stack 中的批量 TPS 统计工具

根目录 README.md 同步增加 GPT-2 Megatron 快速入口。报告、日志、cache、artifacts 和 __pycache__ 均未包含在提交中。

数据与模型对齐

同一份 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 直接加载支持 TP=1、PP=1；数据并行用例支持单卡和 8 卡。

基础用例

Case	精度	GPU / DP	Micro batch	Tokens/step	Sequence	Steps
gpt2_1	FP32	1	4	256	64	10
gpt2_1_bfloat16	BF16 autocast	1	4	256	64	10
gpt2_2	FP32	1	80	5120	64	10
gpt2_2_bfloat16	BF16 autocast	1	80	5120	64	10
gpt2_3	FP32	8	10	5120	64	10
gpt2_3_bfloat16	BF16 autocast	8	10	5120	64	10

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

精度与性能统计

• GPT-2 FP32 loss 最大绝对误差阈值：2e-3
• BF16 loss 最大绝对误差阈值：1e-2
• Megatron iteration 2 的 running-average 日志在比较时恢复为单步 loss
• 性能计时位于 Megatron 标准 train_step 内
• 统计包含数据读取、forward、loss、backward、梯度通信和 optimizer update
• 多卡使用最慢 rank 的耗时
• 共训练 10 步，step 1 作为 warmup，平均 step 2-10
• TPS 不设置 PASS/FAIL 阈值，只报告双方平均 tok/s 和 InfiniTrain/Megatron 百分比

验证结果

使用 InfiniTrain master 日志 scripts/20260904/master_e403ed4/logs/basic 对比：

Case	Loss 最大差异	阈值	结果	InfiniTrain/Megatron TPS
gpt2_1	1.15e-3	2e-3	PASS	186.4%
gpt2_1_bfloat16	2.11e-2	1e-2	FAIL	215.3%
gpt2_2	7.03e-4	2e-3	PASS	88.3%
gpt2_2_bfloat16	9.23e-3	1e-2	PASS	77.9%
gpt2_3	6.99e-4	2e-3	PASS	60.1%
gpt2_3_bfloat16	9.23e-3	1e-2	PASS	53.1%

六个 Megatron 训练用例均运行成功；loss 为 5/6 PASS，TPS 为 6/6 cases compared，无日志缺失或解析错误。

gpt2_1_bfloat16 是 256-token 小 batch，用例对 BF16 舍入和运行版本较敏感。其 Megatron/InfiniTrain 最大差异为 2.11e-2，此前对应 PyTorch/InfiniTrain 对比也达到约 2.30e-2，属于同一数量级。因此保留 1e-2 标准阈值并明确显示 FAIL，不通过放宽阈值隐藏差异。

验证方式

准备数据并运行全部六个 Megatron 用例：

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

只运行选定用例：

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

比较单个用例并输出逐步 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

直接比较 InfiniTrain 标准运行目录和 Megatron 日志目录：

python3 baseline/megatron/scripts/compare_loss.py /path/to/infinitrain/logs/basic baseline/megatron/artifacts/gpt2/logs --include-prefix gpt2_ --include-cases gpt2_1,gpt2_1_bfloat16,gpt2_2,gpt2_2_bfloat16,gpt2_3,gpt2_3_bfloat16 --threshold-fp32 2e-3
python3 baseline/megatron/scripts/compare_tps.py /path/to/infinitrain/logs/basic baseline/megatron/artifacts/gpt2/logs --include-prefix gpt2_ --include-cases gpt2_1,gpt2_1_bfloat16,gpt2_2,gpt2_2_bfloat16,gpt2_3,gpt2_3_bfloat16

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

检查

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

[JYMiracle305]

精度对比

[kilinchange]

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