麻烦补充一下接口设计、约束限制、与其他模块相关性、验收范围、UT设计


Hyper-Parallel 支持 Gated DeltaNet Context Parallel
1. 基本信息
| 项目 | 内容 |
|---|---|
| 作者 | xu-xianliang |
| 相关模块 | trainer、models/qwen3_5、distributed/context_parallel、components/modules、FSDP |
| 相关 Issue / PR | Issue #212;历史实现 PR #1035、#1114;新 Trainer 适配 PR #1362 |
| 适用后端 | PyTorch + Ascend NPU |
2. 背景
Qwen3.5 使用 Full Attention 与 Gated DeltaNet(GDN)混合的层结构。GDN 具有线性复杂度,
但它沿序列维维护 recurrent state。直接切分序列后,每个 CP rank 的计算不仅依赖本地 token,
还依赖前序 rank 的最终状态;三路 causal ShortConv 也需要相邻 rank 的左侧 halo。因此,普通
Full Attention CP 不能直接用于 GDN 层。
旧版实现通过手工调用 apply() 或 parallelize_qwen3_5_cp() 挂载 CP。Hyper-Parallel
训练框架重构后,模型接入统一由 Model Adapter、Sharding Planner 和声明式
plan_overrides 完成,旧挂载接口不再属于当前 Trainer 流程。本 RFC 更新 GDN CP 的接入
方式:用户在训练 YAML 中声明目标模块与 @inner_wrapper,Planner 在 linear_attn 边界
原子替换完整 forward,FSDP、optimizer、checkpoint 继续使用 Trainer 统一流程。
本 RFC 要解决的问题是:在新 Trainer 架构中,为 Hugging Face Qwen3.5 的 GDN 层提供
可训练、可校验、显式失败的 Ulysses 和 State-P2P Context Parallel。
完成后的成功标准:
Qwen3_5ForConditionalGeneration与Qwen3_5ForCausalLM均可被 Model Adapter 识别。- 用户可通过
plan_overrides显式选择ulysses或p2p,以及eager或tritonbackend。 - CP wrapper 覆盖完整 GDN 层,不遗漏 projection、ShortConv、norm、gate 或 output projection。
- CP + FSDP 训练的输出、梯度及多步 loss/grad-norm 轨迹与非 CP reference 对齐。
- 不支持的设备、dtype、shape、依赖或重复挂载在执行前明确报错,不静默回退。
3. 目标和非目标
3.1 目标
- 支持 Qwen3.5 GDN 完整层的 Ulysses CP 与 recurrent-state P2P CP。
- 支持内置纯 PyTorch
eager数学参考后端,以及基于仓库现有 Triton-Ascend GDN 算子的
triton后端。 - 通过 Model Adapter 注册 GDN 参数角色,通过 Sharding Planner 安装 CP executor。
- 支持 CP 与 FSDP 组合训练、non-reentrant activation checkpoint 和多步 AdamW。
- 为 Triton 固定形状路径提供完整 capability 校验;Eager 保留更一般 shape 的验证能力。
- CP size 为 1 时不触发 GDN CP override,保持 Hugging Face 原始模型行为。
3.2 非目标
- 本期不支持 GDN State-AllGather;当前正式模式仅为
ulysses和p2p。 - 本期不支持 packed/varlen、TND、document reset、sequence packing 和跨样本负载均衡。
- 本期不支持推理 cache、增量解码和返回 recurrent cache。
- 本期不支持 TP + CP 同时作用于 GDN 层;检测到该组合时直接报错。
- 本期不迁移完整 Qwen3.5 模型,也不改变 Full Attention、MLP、MoE 或多模态模块的数学实现。
- 本期不融合 causal ShortConv;ShortConv halo 通信后继续执行 PyTorch eager Conv1D。
- 大 CP 下的 parallel scan、分层 P2P、state packet pipeline 和 topology-aware planner 属于后续
性能研究,不进入本期接口与验收范围。
4. 相关实现参考
| 来源 | 做法 | 限制 | 对本 RFC 的影响 |
|---|---|---|---|
| Hugging Face Qwen3.5 | 定义 GDN 层结构、参数名称、Full/GDN 混合层布局 | 不提供分布式 GDN state CP | 保持模型结构与 forward 参数契约一致 |
| Hyper-Parallel 现有 Full Attention CP | 使用 @inner_wrapper 在 attention 边界安装 CP 实现 |
不处理 GDN recurrent state | 复用新 Trainer 的 Planner 与 wrapper 注入机制 |
| Hyper-Parallel GDN Triton-Ascend 算子 | 提供 chunk GDN 前反向融合计算 | 仅支持已验证的 NPU、dtype 和固定 head shape | 作为 triton backend;不满足 capability 时 fail-fast |
| 历史 GDN CP PR | 验证 Ulysses、P2P、affine summary 和反向 state gradient | 使用旧 Trainer 的手工挂载接口 | 复用数学与 executor,重做 Trainer 接入和组合验证 |
5. 对外接口
5.1 训练 YAML
P2P + Triton 示例:
accelerator:
tp_size: 1
cp_size: 4
ep_size: 1
pp_size: 1
plan_overrides:
- match: "*.linear_attn"
when: cp
region_dispatch: false
inner_target: self
inner_wrapper:
_target_: hyper_parallel.models.qwen3_5.adapter.distributed.context_parallel.qwen3_5_gdn_p2p_cp_wrapper
backend: triton
chunk_size: 64
切换为 Ulysses 时仅替换 wrapper:
inner_wrapper:
_target_: hyper_parallel.models.qwen3_5.adapter.distributed.context_parallel.qwen3_5_gdn_ulysses_cp_wrapper
backend: triton
chunk_size: 64
Qwen3.5 是 Full Attention 与 GDN 混合模型。*.linear_attn 只处理 GDN 层;Full Attention
仍需按训练方案配置仓库现有 wrapper,例如 P2P/Ring 路线使用 sdpa_hf_cp_wrapper,纯
Ulysses 路线使用 sdpa_hf_ulysses_cp_wrapper。两类 override 分别匹配不同模块边界。
5.2 配置语义
| 配置项 | 默认值 | 含义 | 合法范围与错误处理 |
|---|---|---|---|
accelerator.cp_size |
1 |
CP mesh 大小 | 大于 1 时 when: cp 生效;等于 1 时不挂 wrapper |
match |
无 | Planner 匹配的模块 FQN | GDN 使用 *.linear_attn |
when |
无 | override 生效条件 | 本特性使用 cp |
region_dispatch |
无 | 是否由框架自动分发 wrapper 内部张量 | 必须为 false,通信和 local tensor 边界由 GDN executor 管理 |
inner_target |
无 | 被替换的目标 | 使用 self,表示替换整个 linear_attn.forward |
inner_wrapper._target_ |
无 | Planner 调用的 @inner_wrapper |
显式选择 Ulysses 或 P2P,不提供 auto 模式 |
backend |
eager |
GDN local core 后端 | eager 或 triton;Triton 不可用时直接报错 |
chunk_size |
64 |
GDN chunk 大小 | Triton 当前只接受 64;Eager 可用于一般 shape 验证 |
5.3 接口设计说明
模式由 wrapper 路径显式编码,backend 由 wrapper 参数显式配置。该设计避免 auto 模式在
依赖、设备或 shape 不满足时静默切换实现,也避免旧版全局 hook 对同一模块重复包装。
@inner_wrapper 不替换模型对象或参数,只向 Planner 返回一个原子的 forward rewrite:
- 校验 CP mesh、TP 组合和重复挂载。
- 为目标 GDN 模块构造 Ulysses 或 P2P executor。
- 将目标模块的完整
forward替换为 executor 调用。 - 写入
_hp_gdn_cp_config标记;第二次挂载同一模块时报RuntimeError。
6. 方案设计
6.1 新 Trainer 总体流程
flowchart TD
A["读取训练 YAML"] --> B["构建 Hugging Face Qwen3.5"]
B --> C["Model Adapter 按 architecture/model_type 注册"]
C --> D["构建 TP/CP/EP/PP DeviceMesh"]
D --> E["Sharding Planner 合并参数规则与 plan_overrides"]
E --> F{"cp_size > 1?"}
F -- "否" --> G["保持 Hugging Face 原始 forward"]
F -- "是" --> H["调用 GDN @inner_wrapper"]
H --> I["原子安装完整层 CP executor"]
I --> J["按 Trainer 配置应用 FSDP"]
J --> K["forward / backward / optimizer / checkpoint"]
Model Adapter 同时注册两个 Hugging Face 身份:
Qwen3_5ForConditionalGeneration/model_type=qwen3_5;Qwen3_5ForCausalLM/model_type=qwen3_5_text。
两者复用同一个 Qwen3.5 provider 和同一组 GDN 参数角色,不需要建立两个模型目录。
6.2 完整 GDN 层边界
wrapper 接管完整 linear_attn.forward,而不只替换 local GDN core:
hidden_states
-> Q/K/V/Beta/Gate projections
-> causal ShortConv halo exchange + local Conv1D
-> q/k L2Norm
-> local GDN prepare / recurrent core
-> Ulysses A2A 或 P2P state propagation
-> output gate + gated RMSNorm
-> output projection
-> attention output
这样可以在一个边界内统一处理 sequence shard、head shard、Conv1D halo 和 recurrent state,
避免只替换核心算子时上下游布局不一致。
6.3 Ulysses CP
Ulysses 将本地 sequence shard 转为完整序列、部分 head:
local Q/K/V/... [B, T/P, H, D]
-> all-to-all: [B, T, H/P, D]
-> local GDN backend 在完整序列上计算
-> reverse all-to-all: [B, T/P, H, D]
优点是复用单卡 GDN 算子,forward/backward 结构直接;代价是交换 token 激活,并要求 K/V
head 数均可被 CP size 整除。
6.4 State-P2P CP
一个 rank 的 local recurrence 可写成 affine map:
H_out = M @ H_in + S
其中 M 表示本地序列对输入 state 的线性变换,S 表示输入 state 为零时由本地 token
产生的状态增量。即使全局初始状态 H0=0,M 仍不可删除:rank 0 的输出是 S0,但
rank 1 需要计算 M1 @ S0 + S1,后续 rank 也都需要用各自的 M 传播非零前缀状态。
Forward:
每个 rank 并行计算 local prepare 与 (M_r, S_r)
rank 0: H_0 = 0
rank r: recv H_r -> H_{r+1} = M_r @ H_r + S_r -> send
每个 rank 使用 H_r 完成本地 GDN output
Backward:
每个 rank 计算 local output gradient 与 gradient-summary
从高 rank 向低 rank 传递 state gradient
dH_r = local_gradient_summary_r + M_r^T @ dH_{r+1}
再由 local GDN backward 生成输入与参数梯度
P2P 保留全部 head、只传递与序列长度无关的 recurrent state,因而不要求 head 数可被 CP
size 整除,也不交换完整 token 激活。代价是 rank 间 state 链具有 O(P) 依赖深度。
6.5 Eager 与 Triton backend
eager 是 Hyper-Parallel 内置的纯 PyTorch 数学参考实现,用于:
- 一般 head shape、FP32 和非 64 整除序列的正确性验证;
- recurrent 与 chunkwise 实现交叉验证;
- CPU UT 和没有 Triton-Ascend 的环境。
triton 调用仓库现有 GDN Triton-Ascend custom ops,用于受支持 NPU shape 的训练性能路径。
backend 在第一次执行 GDN core 时校验依赖与 capability,不满足条件时抛出带具体原因的异常,
不会回退到 Eager。
6.6 参数角色与 FSDP
Model Adapter 为 Planner 补充以下 GDN 参数角色:
| 参数/模块 | Planner 角色 |
|---|---|
in_proj_qkv |
FUSED_QKV |
in_proj_z、in_proj_b、in_proj_a、conv1d |
COLWISE |
A_log、dt_bias |
COLWISE |
out_proj |
ROWWISE |
CP 注入属于模型执行计划,FSDP 由 Trainer 按 mesh 对计划后的模型执行分片。用户不再手工决定
“先 apply CP 还是先 wrap FSDP”;Planner 和 Trainer 固定这一生命周期。GDN executor 不持有
参数副本,参数仍由原模块与 FSDP 管理。
6.7 代码改动点
| 模块 | 改动内容 | 对已有行为的影响 |
|---|---|---|
components/modules/gated_delta_net.py |
Eager reference、backend dispatch、Triton capability 校验 | 仅选择 GDN backend 时生效 |
distributed/context_parallel/gated_delta_net.py |
完整层 Ulysses/P2P executor、state summary 与梯度传播 | 仅被 GDN wrapper 调用 |
models/qwen3_5/adapter/registration.py |
注册两个 HF 身份和 GDN 参数角色 | 不改变其他模型 adapter |
models/qwen3_5/adapter/distributed/context_parallel.py |
两个 @inner_wrapper 与原子 forward rewrite |
仅 when: cp 且匹配 *.linear_attn 时生效 |
trainer |
复用通用 plan_overrides、Planner、FSDP 和 checkpoint 流程 |
不新增 GDN 专用 Trainer 分支 |
tests |
Adapter、YAML、Planner、wrapper、Eager 数学契约 UT | CPU 门禁不导入/执行 Triton kernel |
6.8 方案取舍
| 方案 | 优点 | 缺点 | 本期选择 |
|---|---|---|---|
| Ulysses | 复用 local GDN;无 rank recurrence 链 | A2A token 通信;head 必须可整除 CP | 是 |
| State-P2P | 状态通信量与序列长度无关;保留全部 head | state 链深度为 O(P) |
是 |
| State-AllGather | 无逐 rank state send/recv;可本地合并前缀 | 每 rank 收集全部 summary,通信和显存随 CP 增长 | 否,后续研究 |
手工 apply()/全局 hook |
旧代码迁移量小 | 绕过新 Planner,重复挂载和生命周期难管理 | 否 |
Planner + @inner_wrapper |
与新 Trainer 一致;边界和错误处理明确 | 需要模型 adapter 和显式 YAML | 是 |
7. 组件依赖
| 依赖组件 | 依赖强度 | 当前状态 | 未 ready 时行为 |
|---|---|---|---|
| PyTorch Model Adapter / Sharding Planner | 强依赖 | 已支持 | 无法安装 GDN CP wrapper,直接报配置错误 |
| CP DeviceMesh 与 HCCL | 强依赖 | 已支持 | CP size 大于 1 时初始化失败 |
| FSDP | 组合能力 | 已验证 CP8 + FSDP8 | 可只验证 CP executor,但不构成本 RFC 完整训练验收 |
| Activation checkpoint | 组合能力 | 支持 non-reentrant 路径 | 不影响无 checkpoint 基础功能 |
| Triton-Ascend | Triton backend 强依赖,Eager 不依赖 | 已验证版本从 3.2.1 起 | 选择 Triton 时 fail-fast;普通 import 与 Eager 不受影响 |
| GDN Triton custom ops | Triton backend 强依赖 | 仓库内已有 | 不复制到 model adapter;缺失时 Triton 报错 |
| TP / PP / EP | 本期弱依赖或不涉及 | TP + CP 明确拒绝 | 用户需关闭 TP;PP/EP 不作为本期验收项 |
| MindSpore backend | 不涉及 | 本期未实现 | 不暴露同名 MS 能力 |
8. 约束与兼容性
| 类型 | 当前约束或行为 |
|---|---|
| 模型身份 | 支持 Qwen3_5ForConditionalGeneration 和 Qwen3_5ForCausalLM |
| CP 模式 | ulysses、p2p;无 auto 与正式 all_gather |
| Eager | 支持一般测试 shape、FP32 以及非 chunk 整除序列 |
| Triton dtype | token compute 为 BF16;gate 累计与 recurrent state 为 FP32 |
| Triton shape | Dk=Dv=128、chunk_size=64、local sequence 可被 64 整除 |
| Ulysses head | K/V head 数都必须能被 CP size 整除 |
| P2P head | 不要求 head 数能被 CP size 整除;要求 V heads 可按模型规则映射 K heads |
| Conv1D | 仅 causal、stride=1、depthwise groups 与 channel 契约匹配的训练路径 |
| mask | 支持本地 2D padding mask;因果性由 GDN recurrence 与 causal Conv1D 保证,不构造 S x S mask |
| packed/varlen | 不支持 cu_seq_lens_q、seq_idx、document reset |
| cache/inference | 不支持 recurrent cache 与增量解码 |
| TP + CP | 当前直接抛出 NotImplementedError |
| 重复挂载 | 同一个 GDN 模块第二次 apply wrapper 抛出 RuntimeError |
| CP=1 | when: cp 不生效,不挂 wrapper,保持 Hugging Face 原始路径 |
| 版本行为 | Triton-Ascend 只设最低已验证版本与 capability 检查,不设置未经证明的静态上限 |
9. 验证设计与结果
9.1 CPU UT
当前最小门禁为 5 个 CPU 用例,覆盖:
- 两个 Qwen3.5 HF 身份映射到同一 adapter provider,两个 wrapper 均带
@inner_wrapper标记。 - YAML 的
when: cp在 CP1 不生效,在 CP4 正确生成 GDN override。 - CP4 Planner 能识别
linear_attn完整执行边界并保留 wrapper。 - 原子 forward rewrite、重复挂载拦截和 TP + CP 拦截。
- Eager chunk partition 的 forward/backward 数学一致性。
通过标准为 5 passed。CPU UT 不安装、不导入、不执行 Triton-Ascend kernel;它看护的是
可选依赖隔离和框架契约,而不是 NPU 数值精度。
9.2 NPU 单层与 checkpoint
| 场景 | 检查项 | 通过标准 |
|---|---|---|
| CP4 Ulysses/P2P,Eager/Triton | 完整 GDN output、hidden gradient、参数 gradient | BF16 relative-L2 不高于既定 1e-2 量级,全部 finite |
P2P 非零 h0/dHT |
final state、dH0、全部输入梯度 |
不丢失跨 rank state gradient,不反转 rank 顺序 |
| CP4 non-reentrant checkpoint | output、hidden/参数梯度、通信次序 | 与关闭 checkpoint 同量级;无死锁、消息复用或持续显存增长 |
9.3 新 Trainer 的 CP8 + FSDP8 多步训练
已完成四层混合模型的 200 个 AdamW optimizer step:
- 层布局:3 层 GDN + 1 层 Full Attention;
batch=1,global/local sequence=1024/128;- hidden/intermediate=
2048/6144; - GDN K/V heads=
16/32,head dim=128/128; - Full Attention Q/KV heads=
32/16,head dim=64; - vocab=
4096,BF16,AdamWlr=1e-4、weight_decay=0.01; - 使用确定性的 mmap token stream;reference 为不分片完整模型;candidate 为 CP8 + FSDP8。
| 模式 | 最大 loss abs/rel | 最大 grad-norm abs/rel | 结果 |
|---|---|---|---|
| Ulysses + Triton | 5.5127e-3 / 7.3195e-4 |
5.7570e-2 / 6.1147e-3 |
通过 |
| P2P + Triton | 5.5361e-3 / 6.9334e-4 |
4.9196e-2 / 7.3239e-3 |
通过 |
两种模式均未出现 NaN/Inf、漏梯度、通信死锁或随步数持续扩大的单向漂移。
9.4 性能与显存
性能验证使用固定 local sequence,分别比较同 shape 的 Ulysses 与 P2P;计时需 warmup,取
多次迭代和所有 rank 的最大值。P2P 不设“所有 shape 必须快于 Ulysses”的硬门禁,因为结果
同时受 CP size、local sequence、head 数和互联拓扑影响。验收要求是:
- 新 Trainer 注入不相对相同 executor 的旧路径产生明显额外开销。
- P2P 不交换完整 token 激活,forward 峰值显存应不高于同 shape Ulysses。
- 性能报告必须同时给出 mode、backend、global/local sequence、head shape、warmup/repeat、
forward、backward 和 peak allocated/reserved,不能只报告比例。
历史同 executor 的长序列结果可作为回归参照,但新 Trainer 的正式性能结论应由同环境同轮
A/B 重新生成,不直接沿用旧 Trainer 的数据。
9.5 工程检查
git diff --check、Pythoncompileall、pylint 和 Lizard 门禁通过。- wheel 构建与隔离安装后,GDN module、CP executor、model adapter 和 Triton kernel 均存在。
- 不安装 Triton-Ascend 时可正常 import Hyper-Parallel 并运行 Eager。
- 选择 Triton 但依赖、设备、版本、dtype 或 shape 不满足时,错误信息指出具体 capability。
10. 实现计划与状态
| 阶段 | 内容 | 验证 | 状态 |
|---|---|---|---|
| 1 | 迁移 GDN Eager/Triton backend 与 Ulysses/P2P executor | Eager 数学、NPU 单算子与 state gradient | 已完成 |
| 2 | 新增 Qwen3.5 Model Adapter、参数角色和两个 @inner_wrapper |
Adapter/Planner/YAML CPU UT | 已完成 |
| 3 | 接入新 Trainer 的 plan_overrides,删除对旧手工 apply 流程的依赖 |
CP4 完整层与 checkpoint | 已完成 |
| 4 | CP8 + FSDP8 四层混合模型 200 步训练 | loss、grad norm、finite、漂移 | 已完成 |
| 5 | 补充正式训练 YAML、用户文档和同轮性能/显存报告 | 示例可运行、指标可复现 | 持续完善 |
| 后续 | AllGather、parallel scan、分层/拓扑感知 P2P、packed/varlen、TP + CP | 独立 RFC 与大规模集群验证 | 不属于本期 |
本 RFC 对应的新 Trainer 适配已经由 PR #1362 实现。Issue #212 后续用于记录文档、回归和
本期范围内的验收闭环;大 CP 算法演进和 packed/varlen 支持应拆分独立 Issue,避免扩大当前
接口范围。


Hyper-Parallel 支持 Gated DeltaNet Context Parallel
1. 基本信息
model/distributed/custom_ops/trainer本 RFC 将 Gated DeltaNet(GDN)Context Parallel 的数学语义、通信算法和
本地 GDN 计算后端分层描述:
2. 背景
Qwen3.5 dense decoder 以 3:1 的比例混合 Gated DeltaNet Linear Attention 与 Full
Attention。Hyper-Parallel 原有 Full Attention CP 可以通过 Ulysses 在 sequence shard 和
head shard 之间转换 Q/K/V,但不能直接处理 GDN 的两类跨 rank 因果依赖:
如果每个 rank 只使用零初始状态独立计算本地序列,Conv1D 边界和 GDN 状态都
会错误。这会使 Qwen3.5 混合模型无法在 FSDP+CP 下保持单卡数学语义。
2.1 GDN 计算特征
GDN 不显式构造
[S,S]attention score,而是维护状态矩阵:忽略 chunkwise 并行化细节后,可以将单 token 计算概念化为:
实际 GDN 将 token 按
chunk_size=64分块,通过 gate cumsum、chunk KKT、下三角求解、WY 表示、state scan 和 output kernel 完成前反向。无论如何
分块,rank
r的首状态都必须等于 rankr-1的末状态。2.2 从完整状态扫描到仿射 summary
当本 rank 的 Q/K/V/g/beta 固定后,整段局部序列对输入状态的作用可写成:
其中:
S [B,H,K,V]是H_in=0时本地序列产生的 final state。M [B,H,K,K]描述本地序列对任意输入状态的线性变换。对两个按序执行的区间:
其组合为:
该 compose 满足结合律但不满足交换律。State-P2P 使用本地
(M,S)尽早产生H_out;Summary-AllGather 则收集各 rank(M,S)后按 rank 顺序计算 prefix。3. 目标和非目标
3.1 目标
ContextParallel,GDN 使用LinearAttentionContextParallel。ulysses/p2p/all_gather三种 GDN CP mode。M@H+S和小状态通信。
验证。
3.2 非目标
tp=1。backend=auto,避免环境或 shape 不满足时静默转为 eager。all_gather + triton;AllGather 仅保留 eager 对比路径。纳入当前稳定接口。
4. 相关实现参考
融合 kernel 以源码形式随 Hyper-Parallel wheel 发布,保留原始版权声明并包含
hyper_parallel/platform/torch/custom_ops/gdn/LICENSE。5. 对外接口
5.1 接口定义
CP 样式入口:
LinearAttentionContextParallel( mode="p2p", # "ulysses" | "p2p" | "all_gather" backend="triton", # "eager" | "triton" ).apply(linear_attention_module, cp_mesh)Qwen3.5 模型并行化入口:
parallelize_qwen3_5_cp( model, cp_mesh, ulysses_degree=None, linear_attention_cp_mode="p2p", linear_attention_gdn_backend="triton", )Trainer YAML 配置:
train: accelerator: cp: 4 tp: 1 ulysses_degree: null linear_attention_cp_mode: p2p linear_attention_gdn_backend: tritoncpint1>=1cp维时报错ulysses_degreeOptional[int]NoneNone或cp_sizeNotImplementedErrorlinear_attention_cp_modestr"ulysses"ulysses/p2p/all_gatherlinear_attention_gdn_backendstr"eager"eager/triton5.2 mode/backend 支持矩阵
不提供
auto。用户显式选择triton后,若设备、dtype、shape 或Triton-Ascend 版本不符合契约,立即报错,不会在运行时静默转成 eager。
5.3 使用示例
稳定功能基线:
train: accelerator: cp: 4 tp: 1 linear_attention_cp_mode: ulysses linear_attention_gdn_backend: eager长序列融合 P2P:
train: accelerator: cp: 8 tp: 1 linear_attention_cp_mode: p2p linear_attention_gdn_backend: triton5.4 接口设计说明
mode只选择分布式算法,backend只选择本地 GDN 计算实现。二者正交,可以在不改变 CP 语义的情况下对比 eager/Triton,也可以在同一 Triton backend
上对比 Ulysses/P2P。默认
ulysses+eager保证新增 Triton 依赖不改变已有行为。6. 方案设计
6.1 总体分层
模型层保持 projection、gate、norm 和 out projection 的 Qwen3.5 语义;CP 层管理布局、
halo、collective/P2P 和 distributed autograd;backend 层管理 GDN 前反向;Triton
层不感知 CP mesh 和 rank。
6.2 公共输入和 Conv1D halo
Linear Attention CP 边界输入为 local sequence shard:
若输入为 DTensor,wrapper 在 CP 边界取得 local tensor。各 rank 先执行本地
QKV/Z/A/B projection。
P2P/AllGather 保留 sequence shard,因此 causal Conv1D 需要左侧 halo:
实现使用 differentiable all-to-all-v 传递有效 payload:每个 rank 仅向下一 rank
发送尾部
[B,halo_width,conv_dim],每个非 rank0 只接收前一 rank halo。这不是全量 all-gather,反向由平台可微 collective 处理。Conv 本体保持 eager depthwise Conv1D。
6.3 Ulysses GDN CP
Ulysses 使每个 rank 持有完整 sequence 和部分 heads,所以 Conv 和 recurrent state 都不
跨 rank。其优点是可直接复用完整 local GDN;代价是前反向的大张量 all-to-all,
以及 CP 增大后每 rank GDN heads 减少、单 head 扫描序列变长。
6.4 State-P2P GDN CP
P2P 模式保持每个 rank 拥有全部 GDN heads 和连续 local sequence,只传递 FP32
state:
以 Qwen3.5 GDN core
B=1,H=32,K=V=128为例,单份 state 约为 2 MiB,与sequence length 无关。
6.4.1 eager 仿射 summary
eager P2P 使用 PyTorch 小算子先计算 local prepared tensors 和
(S,M):它是 Triton P2P 的语义基线,也可以在没有 Triton-Ascend 的环境下验证 CP。
6.4.2 Triton 正向流程
各 rank 并行执行不依赖
H_in的 prepare/summary。rank>0在 summary 前预先 postirecv;收到H_in后只在串行链上执行 FP32M@H+S,随后立即isendH_out。本 rank 生成 token states/output 时,下一 rank 已可以处理其状态。isend不是不需要 wait。实现只是将 wait 推迟到本地 output 之后,在 send buffer被释放前必须等待 work 完成。
6.4.3 Triton 反向流程
反向沿 rank 反序传递 state gradient:
state-gradient summary 只是用来尽早生成并发送
dH_in,不是 GDN 的全部反向。Q/K/V/g/beta 的梯度由
backward_state和backward_finish继续计算。GDN本身没有额外的可训练 weight;projection、Conv、norm 和 out projection 参数由外层
PyTorch autograd/FSDP 计算梯度。
6.5 Summary-AllGather
AllGather 取消了 rank-to-rank state receive 链,但每 rank 都需保存所有 summary,并重复
prefix merge。当前只保留 eager 实现用于精度及算法对比,不作为融合后端交付路径。
6.6 GDN backend 组织
Triton 导入是 lazy 的:普通 Hyper-Parallel CPU/MindSpore import 不会导入 GDN Triton
kernel。只有用户显式选择
backend=triton且进入 GDN 路径时才执行版本、NPU、dtype 和 shape 检查。
6.7 代码改动点
core/context_parallel/linear_attention_context_parallel.pymodels/modules/linear_attention.pymodels/qwen3_5/parallelize.pycp=1时不进入新路径platform/torch/custom_ops/gdntrainer/config.pysetup.pycheckpointoptimizer6.8 方案取舍
M@H+S和 state send;保持全 headsbackend=auto7. 组件依赖
DeviceMesh>=3.2.1,<3.38. 约束与兼容性
cp=1、非 Qwen3.5 模型、未选择 Triton 时不增加融合依赖triton-ascend >= 3.2.1, < 3.3,Pythontritonmodule 为 3.2 系列并包含 Ascend backendg为 FP32head_k_dim=head_v_dim=128,chunk_size=64stride=1、depthwisegroups=conv_dim、local sequence 不小于 halo widthwait()前不能释放all_gather+triton直接报错,不回退 eager9. 验证设计与结果
9.1 测试环境与统一口径
主要验证环境:
性能使用
warmup=5, repeat=20。计时边界前后执行 CP barrier 和 device synchronize,并对各 rank elapsed 取
MAX,报告 median/min/p90。显存报告所有 rank 中的peak allocated 最大值。首次 Triton compile/autotune 不计入稳态时间。
9.2 用例分层
<1e-2,全部 finite9.3 单卡 GDN 融合算子
GDN core 实际输入为
B=1,H=32,K=V=128。下表只表示 GDN core,不包含projection、Conv、norm、out projection 或 CP 通信。
序列增长后误差没有系统性放大,全部张量 finite。
eager backward 保存并遍历大量逐 chunk autograd 节点,所以融合 backward 收益随
local chunk 数增长。64K 仅补测性能和显存,未纳入上述长序列精度矩阵。
9.4 CP4/CP8 完整 Linear Attention 层
完整层配置为
hidden=2048,QK/V heads=16/32,K=V=128,BF16,包含 projection、Conv halo、Q/K norm、GDN、RMSNormGated 和 out projection。
CP4/32K 精度:
B=2,global sequence=8192,valid_length=7777的 padding/partial-chunk 用例中,三种 mode的 padding output 均为 exact zero,且输出和梯度通过验收。
性能与显存:
P2P 保持每 rank 32 个 GDN heads。Ulysses CP4/CP8 则将 repeated heads 降为每 rank
8/4 个,同时使单 head sequence 和 chunk 数增长。因此 CP 增大时 Ulysses 的 state
scan kernel 效率下降,P2P 相对收益扩大。
9.5 FSDP+CP 整网精度
整网为 4 层 Qwen3.5,层型为 3 Linear + 1 Full,使用 BF16、AdamW 和非零
学习率。每一步使用不同真实 token 序列。CP4/CP8 都同时开启 FSDP 参数分片。
所有轨迹无 NaN/Inf、无 P2P 顺序错配、无误差持续扩大。Ulysses/P2P 的数值
误差量级接近,没有观察到 P2P state chain 特有的漂移。AllGather eager 额外完成
10 步功能回归。
9.6 整网性能和显存
固定 local sequence=16K:
完整层收益进入整网后被 MLP、Full Attention、embedding、lm_head、loss 和 FSDP
公共开销按 Amdahl 定律稀释。CP4/64K 四层整网仍保留约 9% 收益;CP8
下 Ulysses 的小-head state scan 与布局通信代价更明显,整网收益扩大到约 15.9%。
9.7 Activation checkpoint 交互
full activation checkpoint 会使 GDN custom autograd forward 在正常 forward 执行一次,并在
backward-time layer recompute 再执行一次;custom backward 仍只执行一次。
四组多步轨迹都完整结束,无死锁和数值异常。额外的 deterministic CP4/P2P
checkpoint-off/full 配对实验中,loss、grad norm 和 optimizer trajectory 逐打印位完全一致。
9.8 Level0 profile 和稳定性
CP4/64K Level0 trace 中,P2P 单次
HcclSenddevice self duration 约为0.12--0.14 ms,中间 rank 约9.8 mscommunication 被计算覆盖。状态张量链路传输带宽不是当前主要瓶颈;后续优化应面向 summary/gradient-summary 构造和
backward_finish。每种分布式 mode 连续执行至少 20 次前反向,未观察到 P2P 消息顺序错配、
async buffer 生命周期错误、死锁、NaN/Inf 或显存持续增长。
10. 实现计划
当前稳定提交边界只包含:
实验性 head/V chunk、persistent producer、fused Conv、All-Scan 和
backend=auto不进入当前公共 API。任何后续 kernel 替换都必须保留当前稳定路径作为 A/B 基线,并重新验证
output、final state、
dH0/dHT、全部输入梯度、CP4/CP8 rank-max 性能和peak allocated。