feat: add Qwen3

[JYMiracle305]

feat: add Qwen3

概述

本 PR 为 InfiniTrain 新增 Qwen3-8B dense decoder-only Transformer 支持入口。将 Qwen3-8B 的结构配置映射到现有 nn::TransformerModel 主干，并扩展通用 CausalSelfAttention 以支持 Qwen3 需要的 Q/K RMSNorm 与 half-split RoPE。

本 PR 主要包含四部分：

1. 新增 example/qwen3 示例程序与 CMake target。
2. 新增 Qwen3-8B 默认模型配置。
3. 新增 LLMC v4 FP32 checkpoint loader，支持按 TP/PP/VPP 拓扑加载权重。
4. 扩展通用 Transformer attention：
   • 新增可选 q_norm / k_norm；
   • RoPE 支持 interleaved 与 half-split 两种维度配对布局。

模型结构

本 PR 对应的 Qwen3-8B 配置为：

配置	值
hidden size	4096
layers	36
Q heads	32
KV heads	8
head dim	128
FFN intermediate size	12288
vocab size	151936
max context	40960
norm	RMSNorm, eps=1e-6
position encoding	RoPE, theta=1e6
MLP	SwiGLU
linear bias	无
embedding/lm_head tie	否

模型继续复用 InfiniTrain 已有的：

• pre-RMSNorm residual block；
• causal self-attention；
• GQA；
• SwiGLU MLP；
• final RMSNorm；
• untied lm_head。

Transformer 核心修改

1. Q/K RMSNorm

Qwen3 在 Q/K projection 之后、RoPE 之前，对每个 head 的 head_dim=128 维 Q/K 分别做 RMSNorm。V 不参与该归一化。

计算顺序变为：

c_attn(x)
-> split Q/K/V
-> q_norm(Q), k_norm(K)
-> RoPE(Q, K)
-> GQA repeat K/V
-> scaled dot-product attention

为此新增：

• TransformerConfig::use_qk_norm
• TransformerConfig::qk_norm_eps
• CausalSelfAttention::q_norm_
• CausalSelfAttention::k_norm_
• state dict 名称 attn.q_norm / attn.k_norm

2. half-split RoPE

原有 RoPE 实现使用 interleaved 维度配对：

(dim 0, dim 1), (dim 2, dim 3), ...

Qwen3 / Hugging Face 使用 half-split 配对：

(dim 0, dim D/2), (dim 1, dim D/2+1), ...

因此 ApplyRotaryEmbedding 新增 rotary_interleaved 参数。Qwen3 配置中该值为 false，表示使用 half-split 布局。

如果 RoPE 布局与 checkpoint 训练时的布局不一致，Q/K 的位置旋转会配错维度，导致 attention 结果错误，因此这个适配是本 PR 的关键正确性修改。

Qwen3 example

新增 qwen3 可执行目标，包含：

example/qwen3/config.h
example/qwen3/main.cc
example/qwen3/checkpoint_loader.h
example/qwen3/checkpoint_loader.cc

main.cc 复用现有训练框架，支持：

• CPU / CUDA；
• FP32 / BF16；
• DDP、TP、SP、PP、VPP、ZeRO；
• Adam 与学习率调度；
• LoRA 注入、加载与保存；
• checkpoint resume / save；
• 周期性文本生成。

LLMC checkpoint loader

新增 qwen3::LoadFromLLMC，读取共享 LLMC v4 FP32 权重格式：

• magic：20240804
• version：4
• header：256 * sizeof(int32_t)
• 权重流：连续 FP32 tensor

loader 会：

1. 校验 magic / version；
2. 从 header 恢复模型结构与 RoPE / norm 配置；
3. 按模型结构计算期望文件大小并校验；
4. 根据 TP/PP/VPP 拓扑决定当前 rank 拥有的权重；
5. 将 Qwen3/Hugging Face 权重名映射到 InfiniTrain 内部模块名；
6. 直接填充 TransformerModel::StateDict()。

主要映射关系：

Qwen3 / Hugging Face	InfiniTrain
embed_tokens	transformer.wte
input_layernorm	ln_1
q_proj/k_proj/v_proj	融合 attn.c_attn
q_norm/k_norm	attn.q_norm/k_norm
o_proj	attn.c_proj
post_attention_layernorm	ln_2
gate_proj	mlp.c_fc2
up_proj	mlp.c_fc
down_proj	mlp.c_proj
model.norm	ln_f
lm_head	lm_head

其中：

gate_proj -> c_fc2
up_proj   -> c_fc

是为了匹配现有 MLP::Forward 中 SwiGLU(c_fc2(x), c_fc(x)) 的实现。

实现依据

适配依据包括：

1. Qwen/Qwen3-8B 官方 config.json；
2. Hugging Face Transformers modeling_qwen3.py；
3. Qwen3 checkpoint 权重名与形状；
4. InfiniTrain 现有 Llama3 loader 与 Transformer example；
5. InfiniTrain parallel linear / embedding 的切分语义。

其中：

• 模型规模与超参数来自官方 config.json；
• Q/K RMSNorm 的存在、插入顺序、V 不归一化来自官方 modeling 实现；
• half-split RoPE 来自 Hugging Face rotate_half 实现；
• gate/up/down 映射来自官方 MLP 公式与 InfiniTrain MLP forward 语义；
• LLMC v4 权重顺序属于本仓库使用的转换格式约定。

总结

本 PR 的核心贡献是：

1. 让 InfiniTrain 通用 Transformer 支持 Qwen3 必需的 Q/K RMSNorm；
2. 让 RoPE 支持 half-split 布局；
3. 新增 Qwen3-8B 配置、LLMC v4 loader 和 qwen3 example；
4. 将 Qwen3 权重映射到现有 TP/PP/VPP 并行训练路径。

[JYMiracle305]

qwen3 文件

/data1/shared/InfiniTrain-dev/data/llmc/qwen3/
  ├── qwen3-8b-fp32.llmc                 # Qwen3-8B FP32 权重，由 HF safetensors 转成 LLMC v4
  ├── qwen3_tokenizer.bin                # Qwen3 tokenizer，由 tokenizer.json 生成，用于 GenerateText/Decode
  └── tinyshakespeare/
        ├── tiny_shakespeare.txt           # 原始 Tiny Shakespeare 文本
        ├── tiny_shakespeare_train.bin     # 用 Qwen3 tokenizer 重新编码后的训练 token 流
        └── tiny_shakespeare_val.bin       # 用 Qwen3 tokenizer 重新编码后的验证 token 流

训练示例

tokenizer验证

[kilinchange]

这里 v1/v2/v3 是分别对应给 gpt2/llama3/qwen 使用的 tokenizer format 吗？麻烦确认下，如果是的话在这里加下注释吧

[kilinchange]

这里应该做成一个通用的选项，而不是 qwen3 特有配置，建议使用 bool 类型的 qk_norm 表示是否将当前已有的 norm_type 设置应用到 q 和 k 上，eps 则直接复用 norm_eps。
https://github.com/NVIDIA/Megatron-LM/blob/d737da53c8d0fb2b1958a384312e996796c83a35/megatron/core/transformer/transformer_config.py#L286

[kilinchange]

建议这个配置项也跟 megatron 对齐，直接使用 ffn_hidden_size 表示 ffn 层的 hidden_size，如果没传则默认用 4 * hidden_size。
https://github.com/NVIDIA/Megatron-LM/blob/d737da53c8d0fb2b1958a384312e996796c83a35/megatron/core/transformer/transformer_config.py#L198

[kilinchange]

megatron 里这个配置项默认值是 false，建议对齐默认值。
https://github.com/NVIDIA/Megatron-LM/blob/d737da53c8d0fb2b1958a384312e996796c83a35/megatron/core/transformer/transformer_config.py#L269

[kilinchange]

函数签名尽可能不要使用默认值。

[kilinchange]

qwen3 的指令也需要适配下 #206 的修改，使用 ${DEVICE_BACKEND} 设置 device。

[kilinchange]

这里的默认值与 test_config.json 不一致，建议统一下。

[kilinchange]

建议补充测例：

1. 3d 并行；
2. zero；
3. lr_scheduler；
4. checkpoint；
5. 多进程分布式测例。

[kilinchange]

参考 GPT2/LLaMA3，模型 config 的合法性建议由 config.h 中的 Sanitize 统一校验；这里的单测主要覆盖本次新增的框架能力，不需要重复校验 example 中的具体配置常量，也避免 tests 依赖 example/qwen3/config.h。

[kilinchange]

适配 #206 ，放到上面的 endif() 块内。