RFC: HyperParallel Trainer 新增 GLM-5 系列模型 #2099

需求背景 & 价值

任务地址：https://gitcode.com/mindspore/community/issues/2099
GLM-5 是智谱 AI 于 2026 年 2 月发布的最新 decoder-only 稀疏大语言模型（arXiv:2602.15763），在中文 NLP 生态中占据核心地位。当前 HyperParallel 仅支持 Qwen 系列模型，缺少 GLM 系模型覆盖。

验收标准

要求	标准
训练接入	GLM5 模型在 Trainer 框架中完成 forward / loss / backward / optimizer 完整训练循环
Checkpoint	支持 save_checkpoint / load_checkpoint，Resume 后 loss 连续
集成测试	tests/torch/integration/glm5/ 可运行
配置文档	examples/glm5/ 提供 YAML 配置模板
精度	单卡 100 step 平均 loss 误差 ≤ 5e-3（双卡与单卡对齐）

交付范围

本次交付聚焦训练基础设施接入（Phase 1），使用简化 dense 架构验证注册机制、Trainer 对接、checkpoint 闭环、精度对齐。完整 GLM-5 架构（MoE + MLA + DSA + MTP）在 Phase 2–5 逐步叠加。

GLM-5 完整架构（供参考）：

• 总参 744B，激活 40B（256 专家，top-8 激活），80 层
• MLA（Multi-head Latent Attention）— 576 维 KV 潜变量
• DSA（DeepSeek Sparse Attention）— 200K 上下文稀疏注意力
• MTP（Multi-Token Prediction）— 3 层参数共享，推测解码接受率 2.76

核心价值：

• 补齐 GLM 系模型训练支持，覆盖中文大模型关键模型族
• 遵循 ModelSpec + register_spec 注册模式，零侵入 Trainer 代码
• 复用已有 DSA CP（core/context_parallel/）和 MoE 模块（modules/moe.py），为后续完整架构接入铺垫

---

功能描述

1. 模型架构

GLM-5 采用 MLA + DSA + MoE + MTP 四合一架构：

组件	实现	说明
Attention	MLA（576 维 KV 潜变量）	GQA 的替代方案，KV Cache 体积降低 ~75%
Sparse Attention	DSA（内容感知 Top-K）	200K 长上下文下注意力计算量减少 1.5–2×
Feed-Forward	SwiGLU	dense 层用标准 SwiGLU，MoE 层用 expert 路由
MoE	256 experts，top-8 激活	前几层 dense + 后续 MoE 层
Normalization	RMSNorm	标准 weight * normed
Position Encoding	RoPE	标准 Rotary Embedding
MTP	3 层参数共享	推测解码能力，首期不实现

2. 分阶段交付策略

GLM-5 完整架构复杂度较高，采用分层交付策略——每个 Phase 交付一组可独立训练/验证的组件：

Phase	组件	模型规模	验证方式
Phase 1	Dense GQA（简化架构）	~0.5B，全 dense 层	CPU 训练冒烟 + Ascend 精度对齐
Phase 2	+ MoE	~0.5B 激活，dense + MoE 混合	Ascend EP 精度对齐
Phase 3	+ MLA	~0.5B 激活，MLA attention	Ascend KV Cache 正确性
Phase 4	+ DSA	~0.5B 激活，DSA 注意力	Ascend 长序列推理
Phase 5	+ MTP	~0.5B 激活，3 层 MTP	Ascend 推测解码精度

Phase 1 为最小可用交付：简化架构（dense + GQA + SwiGLU + RMSNorm）即可跑通训练闭环，同时验证注册机制、Trainer 对接、checkpoint 保存恢复。后续 Phase 在此基础逐步叠加真实 GLM-5 组件。

3. 模型注册与发现

遵循现有 ModelSpec + register_spec 注册模式。model.name: glm5 触发 auto-discovery，Universal fields 由 _resolve_overrides() 映射到 GLM5Config。

4. Checkpoint 权重转换

支持两套命名方案（GLM-5 标准布局 + GLM-4 旧版布局），tie 场景自动合成 lm_head.weight。后续 Phase（MoE/MLA/DSA）需扩展键映射逻辑处理 expert 权重、MLA 潜变量参数。

5. 并行策略

• AC：逐层 checkpoint_wrapper
• FSDP：逐层 + root wrap，mixed precision
• TP：通过 _tp_plan 声明（Phase 1 dense 层完全兼容，Phase 2+ 需适配 MoE gate/experts）
• EP（Phase 2+）：复用 modules/moe.py + core/expert_parallel/
• CP（Phase 4+）：复用 core/context_parallel/ 已有的 DSAIndexerContextParallel / DSASparseAttentionContextParallel

---

设计方案

1. 文件组织

hyper_parallel/models/glm5/
├── __init__.py       # register_spec("glm5") + _build() + _resolve_overrides()
├── model.py          # GLM5Config, GLM5ForCausalLM（Phase 1 dense 版本）
├── decoder.py        # GLM5Decoder（Phase 2 加入 MoE 分支）
├── attention.py      # MLA attention（Phase 3）
├── moe.py            # MoE router + expert 层（Phase 2，复用 modules/moe.py）
├── dsa.py            # DSA 索引/边界（Phase 4，对接 core/context_parallel/）
├── mtp.py            # MTP 参数共享（Phase 5）
├── checkpoint.py     # HF safetensors 权重加载 + 键映射
├── parallelize.py    # AC + FSDP + EP/CP 策略
└── state_dict.py     # StateDictAdapter

2. GLM5Config（完整参数集）

@dataclass
class GLM5Config:
    # ── 基础参数（Phase 1 使用） ──
    vocab_size: int = 151936
    hidden_size: int = 1024
    intermediate_size: int = 3072
    num_hidden_layers: int = 24       # 训练用小模型；GLM-5 实际为 80
    num_attention_heads: int = 16
    num_key_value_heads: int = 4      # Phase 1 GQA；Phase 3 MLA 后废弃
    head_dim: int = 64                # Phase 1；MLA 后为 kv_lora_rank
    max_position_embeddings: int = 131072
    rms_norm_eps: float = 1e-6
    rope_theta: float = 500000.0
    tie_word_embeddings: bool = True

    # ── MoE 参数（Phase 2 启用） ──
    num_experts: int = 256
    num_experts_per_tok: int = 8
    num_dense_layers: int = 3         # 前 N 层为 dense，其余为 MoE
    moe_intermediate_size: int = 1024

    # ── MLA 参数（Phase 3 启用） ──
    kv_lora_rank: int = 576           # KV 压缩秩
    qk_rope_head_dim: int = 64        # RoPE 维度（MLA 中 Q/K 分离）
    v_head_dim: int = 128

    # ── DSA 参数（Phase 4 启用） ──
    dsa_topk: int = 2048              # 每 token 选中的 Top-K 历史 token
    dsa_indexer_dim: int = 64

    # ── MTP 参数（Phase 5 启用） ──
    num_mtp_layers: int = 3           # MTP 共享层数

Phase 1 仅使用基础参数。MoE/MLA/DSA/MTP 参数在对应 Phase 启用。

3. Phase 1 模型类层次（简化 dense 版本）

GLM5Config (@dataclass)
GLM5RMSNorm (nn.Module) — 标准 weight * normed
GLM5Decoder (nn.Module):
  ├─ input_layernorm: GLM5RMSNorm
  ├─ self_attn: GroupQueryAttention（Phase 1）/ MLA（Phase 3）
  ├─ post_attention_layernorm: GLM5RMSNorm
  └─ mlp: SwiGLUMLP（Phase 1 dense）/ MoEExperts（Phase 2 MoE）
GLM5TextModel (nn.Module): embed + layers + norm + rotary_emb
GLM5ForCausalLM (nn.Module): model + lm_head + _tp_plan + _cp_modules

4. 前向接口

def forward(self, input_ids, labels=None, position_ids=None, attention_mask=None, **kwargs):
    # Phase 1: embed → dense layers (GQA + SwiGLU) → norm → lm_head → loss
    # Phase 2+: MoE routing 在部分层中替代 SwiGLU
    # Phase 3+: MLA 替代 GQA
    # Phase 4+: DSA 稀疏 mask 叠加到 attention_mask
    return {"loss": loss, "logits": logits}

---

实施计划

Phase 1 — Dense GQA 最小训练闭环（对应验收标准）

目标：满足全部验收标准。使用简化 dense 架构（全层 GQA + SwiGLU + RMSNorm），验证注册机制、Trainer 对接、checkpoint 闭环、精度对齐。

Step	内容	产出
1.1	GLM5Config dataclass（完整参数集，含 MoE/MLA/DSA/MTP 预留）	model.py
1.2	GLM5RMSNorm + GLM5Decoder（dense GQA + SwiGLU）	model.py
1.3	GLM5TextModel + GLM5ForCausalLM	model.py
1.4	__init__.py register_spec + parallelize.py AC/FSDP	__init__.py + parallelize.py
1.5	checkpoint.py + state_dict.py（兼容 GLM-4/GLM-5 布局）	checkpoint.py + state_dict.py
1.6	examples/glm5_dense/train.yaml（参照 qwen3.5 dense 模板）	train.yaml

验证（直接对应验收标准）：

ID	检查项	对应标准
UT-01	Config 默认值 + __post_init__ 校验	训练接入
UT-02	Forward shape + loss.requires_grad	训练接入
UT-03	loss.backward() → 所有参数 grad 非 None	训练接入
UT-04	discover_model_spec + get_spec + parallelize	训练接入
CKPT-01	state_dict save → load → forward logits 一致 (atol=1e-5)	Checkpoint
CKPT-02	tie_word_embeddings data_ptr 验证	Checkpoint
CKPT-03	同 seed 两次 forward loss 完全相同	训练接入
ACC-01	单卡 100 step loss 正常下降，无 NaN/Inf（Ascend A2）	精度
ACC-02	双卡（DP=2）100 step 平均 loss 误差 ≤ 5e-3（Ascend A2）	精度
ACC-03	Save → Load → Resume loss 连续（Ascend A2）	Checkpoint
DOC-01	examples/glm5/train.yaml 可被 HyperTrainerConfig 解析	配置文档
DOC-02	tests/torch/integration/glm5/ 可运行	集成测试

Phase 2 — MoE 架构

目标：将 dense MLP 替换为 MoE（前 num_dense_layers 层 dense + 后续 MoE 层），支持 EP。

Step	内容	说明
2.1	GLM5Decoder 支持 layer_type 调度（dense / moe）	类似 Qwen MoE 的 layer_types
2.2	MoE routing：num_experts=256，topk=8	复用 modules/moe.py 的 MoEExperts
2.3	EP 适配：_ep_modules = ["*.experts"]	复用 core/expert_parallel/
2.4	checkpoint 键映射扩展：处理 expert 权重	checkpoint.py

验证（Ascend A2，EP=2）：

ID	检查项
MOE-01	MoE 层 forward shape 正确
MOE-02	EP=2 loss 与单卡对齐（误差 ≤ 5e-3）
MOE-03	expert 权重 save/load 一致性

Phase 3 — MLA 注意力

目标：用 MLA 替代 Phase 1 的 GQA。MLA 使用 576 维 KV 压缩潜变量替代标准 KV Cache，显存降低 ~75%。

Step	内容	说明
3.1	MLA 类：Q/KV 分离投影 + RoPE 分离 + KV 压缩/解压	新增 attention.py
3.2	GLM5Decoder 支持 attn_type 调度（gqa / mla）	向下兼容 Phase 1
3.3	KV Cache 适配 MLA 格式：潜变量（576 维）vs 标准 KV	更新 generate/kv_cache.py
3.4	checkpoint 键映射：MLA 投影权重命名	checkpoint.py

验证：

ID	检查项
MLA-01	MLA forward 与 GQA forward logits 形状一致
MLA-02	KV Cache 格式正确（潜变量维度）
MLA-03	MLA 权重 save/load 一致性

Phase 4 — DSA 稀疏注意力

目标：集成 DSA，支持 200K 长上下文高效推理。复用 core/context_parallel/ 中已有的 DSAIndexerContextParallel / DSASparseAttentionContextParallel。

Step	内容	说明
4.1	DSA indexer 模块：内容感知 Top-K 选择	新增 dsa.py
4.2	DSA + MLA 联合 attention forward	attention.py
4.3	CP 适配：DSAIndexerContextParallel + mask 构造	复用已有 DSA CP 实现
4.4	长序列验证：128K/200K tokens 推理	Ascend A2

验证（Ascend A2，CP=2）：

ID	检查项
DSA-01	DSA Greedy vs 单卡 dense attention 结果一致
DSA-02	128K tokens Prefill latency + Decode tokens/s
DSA-03	CP=2 生成与单卡一致

Phase 5 — MTP 推测解码

目标：实现 3 层 MTP 参数共享，支持推测解码。

Step	内容	说明
5.1	MTP 模块：3 层共享参数 + 独立 norm	新增 mtp.py
5.2	MTP forward：并行预测 3 个 future token	model.py
5.3	推测解码验证：接受率 ≥ 2.0 tokens/step	Ascend A2

验证：

ID	检查项
MTP-01	MTP forward 输出 4 组 logits（主+3 个 MTP）
MTP-02	推测解码接受率达标
MTP-03	MTP 权重 save/load 一致性

---

对外 API

模型构建（Phase 1）

from hyper_parallel.models.glm5 import GLM5Config, GLM5ForCausalLM

# GLM5Config 完整参数集（Phase 1 仅使用基础参数）
# 默认值：vocab_size=151936, hidden_size=1024, num_hidden_layers=24,
#         num_attention_heads=16, num_key_value_heads=4, head_dim=64

config = GLM5Config(num_hidden_layers=4)  # 小模型快速验证
model = GLM5ForCausalLM(config)
output = model(input_ids, labels=labels)   # {"loss": ..., "logits": ...}

YAML 训练（Phase 1）

# examples/glm5_dense/train.yaml
model:
  name: glm5
  weights_path: null
  tokenizer_path: null
  config_overrides:
    num_hidden_layers: 4

data:
  type: preset_pt
  train_path: /path/to/preset_batches.pt
  max_seq_len: 64

train:
  max_steps: 100
  global_batch_size: 4
  micro_batch_size: 1
  seed: 1234
  backend: torch
  init_device: meta
  accelerator:
    dp_shard: 2
    comm_fusion: true
  optimizer:
    type: adamw
    lr: 1.0e-4
    loss_aggregation: rank_average
  mixed_precision:
    enabled: true
    param_dtype: bfloat16
    reduce_dtype: float32
  gradient_checkpointing:
    activation_checkpoint: full
  checkpoint:
    output_dir: outputs/glm5
    save_steps: 50
  debug:
    deterministic: true

python scripts/train_lm.py --config examples/glm5_dense/train.yaml

Phase 2+ 扩展

MoE/MLA/DSA/MTP 通过 config_overrides 和新增 YAML 字段逐 Phase 启用：

model:
  name: glm5
  config_overrides:
    num_experts: 256
    num_experts_per_tok: 8
    num_dense_layers: 3
    kv_lora_rank: 576        # Phase 3 MLA
    dsa_topk: 2048           # Phase 4 DSA
    num_mtp_layers: 3        # Phase 5 MTP

---

使用约束

• Phase 1 为简化 dense 架构（GQA + SwiGLU），非完整 GLM-5 模型，但可独立训练/验证
• MoE/MLA/DSA/MTP 逐 Phase 叠加，后一 Phase 依赖前一 Phase 的模型骨架
• 仓库已有 DSA CP（dsa_context_parallel.py）和 MoE（modules/moe.py），Phase 2/4 直接复用
• MLA 的 kv_lora_rank=576 与标准 GQA 的 KV Cache 格式不兼容，generate 模块需适配（随 Phase 3）
• TP 切分需保证 num_attention_heads / num_key_value_heads 被 tp_size 整除
• 完整 GLM-5（744B）训练需要大规模集群，本 RFC 聚焦架构实现与精度验证（~0.5B 规模的各 Phase 组件）

---

测试设计

Phase 1 单元测试（CPU）

用例 ID	描述	期望
GLM5-UT-01	Config 默认值 + 扩展字段预留	完整参数集正确
GLM5-UT-02	Forward shape (bsz=2, seq=8)	logits (2,8,V)，loss.requires_grad
GLM5-UT-03	loss.backward() 所有参数 grad	非 None
GLM5-UT-04	Forward + attention_mask	shape 一致
GLM5-UT-05	discover_model_spec + get_spec	注册成功
GLM5-UT-06	parallelize_glm5 不抛异常	AC+FSDP 正常

Phase 1 Checkpoint 测试（CPU）

用例 ID	描述	期望
GLM5-CKPT-01	save → load → forward	logits 一致 (atol=1e-5)
GLM5-CKPT-02	tie_word_embeddings data_ptr	共享存储
GLM5-CKPT-03	同 seed 两次 forward loss	完全相同

Phase 1 精度验证（Ascend A2）

用例 ID	描述	期望
GLM5-ACC-01	单卡 100 step loss 下降	无 NaN/Inf
GLM5-ACC-02	双卡（DP=2）vs 单卡 loss 对齐	均值误差 ≤ 5e-3
GLM5-ACC-03	Save → Load → Resume	loss 连续

Phase 2 分布式测试（Ascend A2）

用例 ID	描述	期望
GLM5-EP-01	EP=2 MoE forward shape 正确	与单卡一致
GLM5-EP-02	EP=2 loss 对齐	误差 ≤ 5e-3
GLM5-EP-03	expert 权重 save/load 一致	atol=1e-5

回归测试

• Qwen3.5 / Qwen3.5-MoE / Qwen3-VL-MoE 模型注册与训练不受影响
• tests/torch/integration/llamafactory/ 全部通过

---

规格 & 约束

• Phase 1（验收交付）：满足全部验收标准——训练闭环、checkpoint 保存恢复、集成测试、配置文档、100 step 平均 loss 误差 ≤ 5e-3
• Phase 2–5（后续扩展）：MoE / MLA / DSA / MTP，每 Phase 独立可训练/验证
• 环境：Python 3.10 / PyTorch 2.6 / MindSpore>=2.8 / CANN 8.5
• 硬件：Phase 1 需 2×Ascend A2
• 硬件：Phase 1–3 需 2×Ascend A2；Phase 4–5 长序列测试可能需更多卡

---

参考

• GLM-5 论文：https://arxiv.org/abs/2602.15763
• 现有模型实现：hyper_parallel/models/qwen3_5/、hyper_parallel/models/qwen3_5_moe/
• DSA CP 实现：hyper_parallel/core/context_parallel/dsa_context_parallel.py
• MoE 模块：hyper_parallel/models/modules/moe.py
• 模型注册规范：hyper_parallel/models/spec/
• Trainer 框架：hyper_parallel/trainer/base.py
• SIG 仓库：https://atomgit.com/mindspore/community/tree/master/sigs/parallel_training_system