【RFC】HyperParallel Trainer新增GLM5系列模型 #2099

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

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 182
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/182

Start with the Phase 1 files listed in the issue: hyper_parallel/models/glm5/model.py, init.py, parallelize.py, checkpoint.py, state_dict.py, and examples/glm5_dense/train.yaml. Read the existing Qwen model and Trainer registration patterns before running the listed UT, CKPT, ACC, and DOC checks. Done requires the dense training loop, checkpoint resume, integration tests, configuration template, and stated single- and dual-card loss criteria.