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 并行训练路径。

[Chamberlain0w0]

这里也应该手动把 use_qk_norm = true 加进去，保证不读 llmc path 分支的正确性

[Chamberlain0w0]

感觉没必要额外存两个 q_norm_ 和 k_norm_ 到 class 里，可以正常就按 submodule 注册，然后后面用局部变量获取

[Chamberlain0w0]

同上，感觉没有额外存一个类内成员的必要

[Chamberlain0w0]

这个 tokenizer 可能也得适配下。目前 tokenizer.cc 里面可能是硬编码了 gpt2/llama3 的一些 eos 的值