已关闭
RFC: HyperParallel Linear Attention Context Parallel 架构 #212
xu-xianliang创建于  6月14日关闭于  29 天前
xu-xianliang
xu-xianliang
6月14日 创建

Hyper-Parallel 支持 Gated DeltaNet Context Parallel

1. 基本信息

项目 内容
作者 xu-xianliang
相关模块 model / distributed / custom_ops / trainer
小算子 CP PR Hyper-Parallel PR 1035
融合后端 PR Hyper-Parallel PR 1114
相关 issue / SR / AR 待补充
适用后端 PyTorch
目标硬件 Ascend NPU

本 RFC 将 Gated DeltaNet(GDN)Context Parallel 的数学语义、通信算法和
本地 GDN 计算后端分层描述:

  1. eager/小算子后端是可移植的语义基线。
  2. Triton-Ascend 融合后端是显式开启的性能能力。
  3. Ulysses、State-P2P 和 Summary-AllGather 是 CP 算法,不是 GDN backend。

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 因果依赖:

  1. causal depthwise Conv1D 需要当前 rank 左侧的少量 token halo。
  2. GDN recurrent state 的当前 rank 初始值依赖前一 rank 的 final state。

如果每个 rank 只使用零初始状态独立计算本地序列,Conv1D 边界和 GDN 状态都
会错误。这会使 Qwen3.5 混合模型无法在 FSDP+CP 下保持单卡数学语义。

2.1 GDN 计算特征

GDN 不显式构造 [S,S] attention score,而是维护状态矩阵:

H_t: [B, H, K, V]

忽略 chunkwise 并行化细节后,可以将单 token 计算概念化为:

residual_t = beta_t * (v_t - k_t @ H_(t-1))
H_t        = exp(g_t) * H_(t-1) + outer(k_t, residual_t)
o_t        = q_t @ H_t

实际 GDN 将 token 按 chunk_size=64 分块,通过 gate cumsum、chunk KKT、
下三角求解、WY 表示、state scan 和 output kernel 完成前反向。无论如何
分块,rank r 的首状态都必须等于 rank r-1 的末状态。

2.2 从完整状态扫描到仿射 summary

当本 rank 的 Q/K/V/g/beta 固定后,整段局部序列对输入状态的作用可写成:

H_out = M @ H_in + S

其中:

  • S [B,H,K,V] 是 H_in=0 时本地序列产生的 final state。
  • M [B,H,K,K] 描述本地序列对任意输入状态的线性变换。

对两个按序执行的区间:

H_mid = M0 @ H_in  + S0
H_out = M1 @ H_mid + S1

其组合为:

M = M1 @ M0
S = M1 @ S0 + S1

该 compose 满足结合律但不满足交换律。State-P2P 使用本地 (M,S) 尽早产生
H_out;Summary-AllGather 则收集各 rank (M,S) 后按 rank 顺序计算 prefix。

本 RFC 要解决的问题:为 Hyper-Parallel 补齐 Gated DeltaNet Context Parallel,
并提供显式可选的 Triton-Ascend 融合 GDN 后端。

成功标准:Qwen3.5 Full/Linear Attention 混合模型可在 PyTorch/Ascend NPU
上完成 FSDP+CP4/CP8 多步训练;输出、梯度、loss 和 grad norm 与单卡
基线对齐;融合 P2P 在目标长序列下性能不低于 Ulysses。

3. 目标和非目标

3.1 目标

  1. 支持 Qwen3.5 Full Attention 和 GDN 混合层使用同一一维 CP mesh。
  2. Full Attention 继续复用已有 ContextParallel,GDN 使用
    LinearAttentionContextParallel。
  3. 支持 ulysses / p2p / all_gather 三种 GDN CP mode。
  4. 使用 eager 后端建立无 Triton 依赖的功能和精度基线。
  5. Ulysses 和 P2P 支持显式 Triton-Ascend GDN backend。
  6. State-P2P 正向只沿 sequence rank 传递 final state,反向只传递 state gradient。
  7. 通过仿射 summary 将 P2P 串行链从完整 GDN state scan 缩短为 M@H+S
    和小状态通信。
  8. 正确处理 causal Conv1D halo、DTensor/local tensor 边界和 FSDP 参数分片。
  9. 不支持的 mode/backend/device/dtype/shape 组合必须直接报错,不静默回退。
  10. 在 CP4/CP8 下完成单算子、完整 GDN 层、混合整网和 activation checkpoint
    验证。

3.2 非目标

  1. 本期不支持 MindSpore GDN CP。
  2. 本期不支持 Linear Attention TP+CP 组合;开启 Linear Attention CP 时要求
    tp=1。
  3. 本期不验收 PP/EP 与 GDN CP 的组合。
  4. 本期不支持 inference state cache、decode 和增量生成。
  5. 本期不支持 packed/varlen GDN CP。
  6. 不提供 backend=auto,避免环境或 shape 不满足时静默转为 eager。
  7. 不提供 all_gather + triton;AllGather 仅保留 eager 对比路径。
  8. 不将 fused Conv1D、head/V chunk、Ulysses×P2P、All-Scan 和 prepare-summary 融合
    纳入当前稳定接口。
  9. 不引入 MindSpeed-MM 等框架级运行时依赖。

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: triton
入参 / 配置项 类型 默认值 是否必填 含义 合法范围 错误处理
cp int 1 否 Context Parallel degree >=1 mesh 中没有 cp 维时报错
ulysses_degree Optional[int] None 否 Full Attention Ulysses degree 当前仅支持 None 或 cp_size 非 pure Ulysses 报 NotImplementedError
linear_attention_cp_mode str "ulysses" 否 GDN CP 通信算法 ulysses/p2p/all_gather 未知 mode 报错
linear_attention_gdn_backend str "eager" 否 本地 GDN 计算后端 eager/triton 未知 backend 报错

5.2 mode/backend 支持矩阵

CP mode eager Triton-Ascend
Ulysses 支持,默认基线 支持,完整 local GDN fwd/bwd
State-P2P 支持,PyTorch summary 支持,融合 GDN + state-summary
Summary-AllGather 支持,对比模式 不支持,直接报错

不提供 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: triton

5.4 接口设计说明

mode 只选择分布式算法,backend 只选择本地 GDN 计算实现。二者正交,
可以在不改变 CP 语义的情况下对比 eager/Triton,也可以在同一 Triton backend
上对比 Ulysses/P2P。默认 ulysses+eager 保证新增 Triton 依赖不改变已有行为。

6. 方案设计

6.1 总体分层

Qwen3.5 model / parallelize
        |
        v
LinearAttentionContextParallel(mode, backend)
        |
        +-- Ulysses communication/layout
        +-- State-P2P communication/autograd
        +-- Summary-AllGather communication/prefix
        |
        v
GDN backend contract
        |
        +-- eager PyTorch reference
        +-- Triton-Ascend fused backend
        |
        v
platform.torch.custom_ops.gdn

模型层保持 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:

hidden_states: [B, S/P, hidden]

若输入为 DTensor,wrapper 在 CP 边界取得 local tensor。各 rank 先执行本地
QKV/Z/A/B projection。

P2P/AllGather 保留 sequence shard,因此 causal Conv1D 需要左侧 halo:

halo_width = (kernel_size - 1) * dilation

实现使用 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

local sequence shard
  -> local projection
  -> differentiable A2A: sequence shard -> head shard
  -> full sequence + reduced heads
  -> local causal Conv1D
  -> local GDN (eager or Triton)
  -> inverse differentiable A2A: head shard -> sequence shard
  -> gated norm + out projection

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:

state: [B, H, K, V]

以 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):

all ranks: prepare -> local summary (S, M)

rank chain:
recv H_in -> H_out = M @ H_in + S -> isend H_out

local output:
prepared + H_in -> token states -> token output

它是 Triton P2P 的语义基线,也可以在没有 Triton-Ascend 的环境下验证 CP。

6.4.2 Triton 正向流程

forward_prepare
  -> Q/K norm, gate cumsum, chunk A, W, U

forward_state_summary
  -> S, M

forward_summary_apply
  -> H_out = M @ H_in + S

forward_apply_state
  -> per-chunk states h, v_new

forward_output
  -> local token output

各 rank 并行执行不依赖 H_in 的 prepare/summary。rank>0 在 summary 前预先 post
irecv;收到 H_in 后只在串行链上执行 FP32 M@H+S,随后立即 isend
H_out。本 rank 生成 token states/output 时,下一 rank 已可以处理其状态。

time ---> |                                                                                 |
------------------------------------------------------------------------------------------------
Rank0     | PS0 |A0|send H1| O0                                                             |
Rank1     | PS1 |wait H1 |recv|A1|send H2| O1                                                |
Rank2     | PS2 |                 wait H2 |recv|A2|send H3| O2                               |
Rank3     | PS3 |                                  wait H3 |recv|A3| O3                      |

PS: parallel prepare + state summary
A : state-summary apply, H_out=M@H_in+S
O : local state scan + token output

isend 不是不需要 wait。实现只是将 wait 推迟到本地 output 之后,在 send buffer
被释放前必须等待 work 完成。

6.4.3 Triton 反向流程

反向沿 rank 反序传递 state gradient:

backward_prepare(saved, dO)
  -> W, h, v_new, dv and local backward context

backward_state_gradient_summary
  -> dS_local, M

backward_summary_apply
  -> dH_in = M^T @ dH_out + dS_local

backward_state + backward_finish
  -> dQ, dK, dV, dg, dbeta

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

all ranks: prepare -> local (S, M)
  -> pack S/M
  -> differentiable all_gather
  -> rank r 按序合并 rank 0..r-1 summary
  -> 得到本 rank H_in
  -> local state scan + output

AllGather 取消了 rank-to-rank state receive 链,但每 rank 都需保存所有 summary,并重复
prefix merge。当前只保留 eager 实现用于精度及算法对比,不作为融合后端交付路径。

6.6 GDN backend 组织

hyper_parallel/models/modules/linear_attention.py
  -> eager implementation
  -> explicit backend dispatch + capability check

hyper_parallel/platform/torch/custom_ops/gdn/
  -> chunk_gated_delta_rule.py       # Hyper autograd/operator contract
  -> state_summary.py                # P2P summary/apply/gradient-summary
  -> triton/*.py                     # chunkwise Triton-Ascend kernels
  -> LICENSE

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.py Ulysses/P2P/AllGather wrapper、Conv halo、state-summary 通信与 autograd 否,只有显式开启 GDN CP 时生效
models/modules/linear_attention.py eager GDN oracle、Triton capability 和 dispatch 默认 eager,不改变已有行为
models/qwen3_5/parallelize.py 按 layer type 分别挂载 Full/Linear Attention CP cp=1 时不进入新路径
platform/torch/custom_ops/gdn GDN Triton-Ascend 前反向和 P2P summary kernel 新增可选能力
trainer/config.py 新增 CP mode 和 GDN backend 配置 默认值保持 eager/Ulysses
setup.py wheel 打包 GDN LICENSE 否
checkpoint 无公共接口改动 不涉及
optimizer 无改动 不涉及

6.8 方案取舍

方案 优点 缺点 是否选择 原因
Ulysses 直接复用完整 local GDN,没有 rank state 链 大张量 A2A;CP 增大后 head 并行度降低 是,默认 语义简单、稳定基线
原始完整 GDN P2P 只通信小状态 完整 local GDN 位于 rank 串行链 否 CP 增大时等待累积
Affine-summary P2P 串行链只留 M@H+S 和 state send;保持全 heads 额外计算 summary,依赖定制前反向 是,性能路径 目前长序列/大 CP 下的最优稳定 P2P
Summary-AllGather 无链式 state recv 收集全部 summary、重复 prefix merge 是,eager 对比 用于数学和通信方案对比
backend=auto 用户配置简单 容易因环境/shape 静默改变性能语义 否 必须 fail-fast
Fused Conv1D 可进一步减少 Conv 时间和 activation 导入顺序、shape 和算子审查风险 否,后续 不与 GDN summary 同时扩大提交面
Head/V wavefront、All-Scan 大 CP 下可继续降低状态链 小 kernel 效率和通信原语复杂 否,研究项 需要独立 RFC 和 CP16/32 证据

7. 组件依赖

依赖组件 强依赖 / 弱依赖 当前状态 未 ready 时本期能力
PyTorch DeviceMesh 强依赖 已有 无 GDN CP
HCCL collective/P2P 强依赖 已有 无分布式能力
differentiable all-to-all Ulysses/Conv halo 强依赖 平台层已有 Ulysses 和 halo 路径不可用
Triton-Ascend >=3.2.1,<3.3 融合 backend 强依赖 已验证 eager 三种 CP mode 仍可用
Ascend NPU 融合 backend 强依赖 已验证 只可运行 eager/非 NPU 可执行项
FSDP 目标组合强依赖 CP4/CP8 已验证 仍可做 module CP,无法验收整网目标
activation checkpoint 弱依赖 CP4/CP8 full 已验证 不影响 checkpoint-off 能力
MindSpore 本期不涉及 未实现 PyTorch 能力不受影响
完整融合能力需要:PyTorch + Ascend NPU + HCCL + DeviceMesh +
triton-ascend 3.2.x。

最小可交付能力是:eager Ulysses/P2P/AllGather GDN CP,无 Triton 运行时依赖。

8. 约束与兼容性

类型 内容
已有行为 cp=1、非 Qwen3.5 模型、未选择 Triton 时不增加融合依赖
后端范围 GDN CP 当前只支持 PyTorch;Triton 只支持 Ascend NPU
Triton 版本 triton-ascend >= 3.2.1, < 3.3,Python triton module 为 3.2 系列并包含 Ascend backend
Triton dtype Q/K/V/beta 为 BF16,gate g 为 FP32
Triton shape head_k_dim=head_v_dim=128,chunk_size=64
P2P Triton sequence 每 rank local sequence 必须被 64 整除
Ulysses 相关 head/projection 维必须满足 CP 整除契约;Full Attention 仅支持 pure Ulysses
Conv stride=1、depthwise groups=conv_dim、local sequence 不小于 halo width
P2P 顺序 所有 rank 必须使用对称 tag/order;async send buffer 在 wait() 前不能释放
AllGather all_gather+triton 直接报错,不回退 eager
组合约束 Linear Attention TP+CP、packed/varlen、decode/cache、PP/EP 不在当前验收范围
短序列 P2P summary 存在固定成本,不保证短序列一定快于 Ulysses

9. 验证设计与结果

9.1 测试环境与统一口径

主要验证环境:

Hardware:       Ascend 910B3
CANN:           9.1.0-beta.3
PyTorch:        2.10.0
torch-npu:      2.10.0
triton-ascend:  3.2.1
Compute dtype:  BF16
GDN gate/state: FP32

性能使用 warmup=5, repeat=20。计时边界前后执行 CP barrier 和 device synchronize,
并对各 rank elapsed 取 MAX,报告 median/min/p90。显存报告所有 rank 中的
peak allocated 最大值。首次 Triton compile/autotune 不计入稳态时间。

9.2 用例分层

用例级别 覆盖内容 通过标准
CPU UT mode/backend 参数、unsupported combination、QKV slice、summary pack/merge 及 autograd 全部通过,不导入 Triton kernel
NPU 算子 eager/Triton,8K/16K/32K,随机 H0/dO/dHT,全部输入梯度 rel-L2 <1e-2,全部 finite
CP 完整层 CP4/CP8,Ulysses/P2P/AllGather,padding/partial chunk output/input-grad/parameter-grad/grad norm 对齐,无 hang
整网 Level1 3 Linear + 1 Full,FSDP+CP4/CP8,非零学习率 100/200 step loss 和 grad norm 无持续漂移
checkpoint full activation checkpoint,CP4/CP8 重放时 P2P 顺序稳定,精度不变
性能/显存 CP4/CP8 长序列,rank-max 报告稳态时间、tokens/s 和 peak allocated

9.3 单卡 GDN 融合算子

GDN core 实际输入为 B=1,H=32,K=V=128。下表只表示 GDN core,不包含
projection、Conv、norm、out projection 或 CP 通信。

Sequence 最差 output rel-L2 final-state rel-L2 dH0 rel-L2 最差 input-grad rel-L2
8K 4.3701e-3 3.3544e-3 1.5870e-3 5.6839e-3
16K 4.3707e-3 3.3464e-3 1.5942e-3 5.8559e-3
32K 4.3712e-3 3.3549e-3 1.5840e-3 5.6893e-3

序列增长后误差没有系统性放大,全部张量 finite。

Sequence Triton Forward speedup Triton Backward speedup Triton Fwd+Bwd speedup Incremental peak reduction
8K 7.98x 13.56x 12.65x 3.83x
16K 7.90x 26.40x 22.37x 3.84x
32K 10.09x 50.23x 40.08x 3.84x
64K 10.71x 95.31x 74.17x 3.84x

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 精度:

Mode Output rel-L2 Input-grad rel-L2 Parameter-grad rel-L2 Grad-norm rel
Ulysses Triton 4.4771e-3 5.8584e-3 5.0756e-3 3.7054e-5
P2P Triton 4.4771e-3 5.8587e-3 5.0757e-3 3.7163e-5
AllGather eager 1.3481e-4 3.5978e-4 2.3472e-3 1.0898e-6

B=2,global sequence=8192,valid_length=7777 的 padding/partial-chunk 用例中,三种 mode
的 padding output 均为 exact zero,且输出和梯度通过验收。

性能与显存:

Topology Global/local seq Ulysses Fwd+Bwd P2P Fwd+Bwd P2P time reduction Peak allocated difference
CP4 32K/8K 76.884 ms 65.807 ms 14.4% P2P +12.0 MiB
CP4 64K/16K 160.283 ms 132.973 ms 17.0% P2P +12.0 MiB
CP4 128K/32K 311.687 ms 281.569 ms 9.7% P2P +12.0 MiB
CP8 64K/8K 98.427 ms 67.336 ms 31.6% P2P -52.1 MiB
CP8 128K/16K 279.313 ms 151.257 ms 45.8% P2P -116.2 MiB

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 参数分片。

Topology Mode Steps Max loss abs/rel Max grad-norm abs/rel
CP4 + FSDP4 Ulysses 100 2.6097e-3 / 3.6715e-4 1.7036e-2 / 3.7480e-3
CP4 + FSDP4 P2P 100 2.6107e-3 / 3.6729e-4 1.8070e-2 / 3.9756e-3
CP8 + FSDP8 Ulysses 200 2.6898e-3 / 3.7746e-4 1.7224e-2 / 3.7883e-3
CP8 + FSDP8 P2P 200 2.6526e-3 / 3.7223e-4 1.7531e-2 / 3.8563e-3

所有轨迹无 NaN/Inf、无 P2P 顺序错配、无误差持续扩大。Ulysses/P2P 的数值
误差量级接近,没有观察到 P2P state chain 特有的漂移。AllGather eager 额外完成
10 步功能回归。

9.6 整网性能和显存

固定 local sequence=16K:

Topology Layers Ulysses Fwd+Bwd P2P Fwd+Bwd P2P time reduction Throughput increase Peak allocated reduction
CP4/FSDP4, global 64K 4 854.047 ms 778.644 ms 8.83% 9.68% 372.0 MiB / 2.33%
CP4/FSDP4, global 64K 8 1686.321 ms 1529.381 ms 9.31% 10.26% 744.1 MiB / 2.46%
CP8/FSDP8, global 128K 4 1151.401 ms 967.863 ms 15.94% 18.96% 372.0 MiB / 2.38%
CP8/FSDP8, global 128K 8 2276.529 ms 1915.960 ms 15.84% 18.82% 744.0 MiB / 2.47%

完整层收益进入整网后被 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 仍只执行一次。

Topology Checkpoint Ulysses Fwd+Bwd P2P Fwd+Bwd P2P reduction Ulysses/P2P peak
CP4/FSDP4, 64K/16K off 854.047 ms 778.644 ms 8.83% 15986.9 / 15614.9 MiB
CP4/FSDP4, 64K/16K full 1082.048 ms 990.121 ms 8.50% 5602.3 / 5478.3 MiB
CP8/FSDP8, 128K/16K off 1151.401 ms 967.863 ms 15.94% 15645.2 / 15273.2 MiB
CP8/FSDP8, 128K/16K full 1481.585 ms 1248.249 ms 15.75% 5490.1 / 5366.1 MiB

四组多步轨迹都完整结束,无死锁和数值异常。额外的 deterministic CP4/P2P
checkpoint-off/full 配对实验中,loss、grad norm 和 optimizer trajectory 逐打印位完全一致。

9.8 Level0 profile 和稳定性

CP4/64K Level0 trace 中,P2P 单次 HcclSend device self duration 约为
0.12--0.14 ms,中间 rank 约 9.8 ms communication 被计算覆盖。状态张量链路传输
带宽不是当前主要瓶颈;后续优化应面向 summary/gradient-summary 构造和
backward_finish。

每种分布式 mode 连续执行至少 20 次前反向,未观察到 P2P 消息顺序错配、
async buffer 生命周期错误、死锁、NaN/Inf 或显存持续增长。

10. 实现计划

阶段 / PR 内容 依赖 验证 状态
PR 1035 eager Ulysses/P2P/AllGather、Conv halo A2AV、Qwen3.5 混合层接入 已有 Hyper CP/FSDP CPU UT + NPU 完整层 + 多步精度 已完成并合入
PR 1114 Triton local GDN、P2P state-summary 正反向、backend fail-fast、wheel/LICENSE PR 1035,triton-ascend 3.2.x 单算子、CP4/CP8、100/200 步、profile、checkpoint 已完成实现和验证
Follow-up 1 将 forward summary 融入 W/U prepare,减少 HBM 回写/重读 稳定 P2P Triton output/state/全梯度 + CP4/CP8 A/B 规划中
Follow-up 2 将 gradient-summary 融入 backward prepare,优化长 local sequence backward Follow-up 1 可独立 dH0/dHT/全梯度 + Level0 规划中
Follow-up 3 segmented summary / parallel prefix scan 新 kernel CP8/CP16 扩展、peak memory 研究项
Follow-up 4 Fused Conv1D、AllGather Triton、Ulysses×P2P、All-Scan 独立算子和通信设计 独立 RFC/PR 不阻塞当前交付

当前稳定提交边界只包含:

mode = ulysses | p2p | all_gather
backend = eager | triton

ulysses + triton: 完整 local fused GDN
p2p + triton:     fused GDN + fused state-summary
all_gather:       eager only
conv:             eager only

实验性 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。

likedislike
xu-xianliangxu-xianliang
6月14日 修改了issue 的描述
xu-xianliangxu-xianliang
6月15日 修改了issue 的描述
missyuanzi
8月5日 评论:

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

likedislike
xu-xianliangxu-xianliang
8月11日 修改了issue 的描述
xu-xianliangxu-xianliang
8月14日 修改了issue 的描述
xu-xianliangxu-xianliang
9月9日 关联了pull request:feat(torch): restore GDN context parallel
xu-xianliang
xu-xianliang
9月10日 评论:

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。

完成后的成功标准:

  1. Qwen3_5ForConditionalGeneration 与 Qwen3_5ForCausalLM 均可被 Model Adapter 识别。
  2. 用户可通过 plan_overrides 显式选择 ulysses 或 p2p,以及 eager 或 triton backend。
  3. CP wrapper 覆盖完整 GDN 层,不遗漏 projection、ShortConv、norm、gate 或 output projection。
  4. CP + FSDP 训练的输出、梯度及多步 loss/grad-norm 轨迹与非 CP reference 对齐。
  5. 不支持的设备、dtype、shape、依赖或重复挂载在执行前明确报错,不静默回退。

3. 目标和非目标

3.1 目标

  1. 支持 Qwen3.5 GDN 完整层的 Ulysses CP 与 recurrent-state P2P CP。
  2. 支持内置纯 PyTorch eager 数学参考后端,以及基于仓库现有 Triton-Ascend GDN 算子的
    triton 后端。
  3. 通过 Model Adapter 注册 GDN 参数角色,通过 Sharding Planner 安装 CP executor。
  4. 支持 CP 与 FSDP 组合训练、non-reentrant activation checkpoint 和多步 AdamW。
  5. 为 Triton 固定形状路径提供完整 capability 校验;Eager 保留更一般 shape 的验证能力。
  6. CP size 为 1 时不触发 GDN CP override,保持 Hugging Face 原始模型行为。

3.2 非目标

  1. 本期不支持 GDN State-AllGather;当前正式模式仅为 ulysses 和 p2p。
  2. 本期不支持 packed/varlen、TND、document reset、sequence packing 和跨样本负载均衡。
  3. 本期不支持推理 cache、增量解码和返回 recurrent cache。
  4. 本期不支持 TP + CP 同时作用于 GDN 层;检测到该组合时直接报错。
  5. 本期不迁移完整 Qwen3.5 模型,也不改变 Full Attention、MLP、MoE 或多模态模块的数学实现。
  6. 本期不融合 causal ShortConv;ShortConv halo 通信后继续执行 PyTorch eager Conv1D。
  7. 大 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:

  1. 校验 CP mesh、TP 组合和重复挂载。
  2. 为目标 GDN 模块构造 Ulysses 或 P2P executor。
  3. 将目标模块的完整 forward 替换为 executor 调用。
  4. 写入 _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 用例,覆盖:

  1. 两个 Qwen3.5 HF 身份映射到同一 adapter provider,两个 wrapper 均带 @inner_wrapper 标记。
  2. YAML 的 when: cp 在 CP1 不生效,在 CP4 正确生成 GDN override。
  3. CP4 Planner 能识别 linear_attn 完整执行边界并保留 wrapper。
  4. 原子 forward rewrite、重复挂载拦截和 TP + CP 拦截。
  5. 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,AdamW lr=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 数和互联拓扑影响。验收要求是:

  1. 新 Trainer 注入不相对相同 executor 的旧路径产生明显额外开销。
  2. P2P 不交换完整 token 激活,forward 峰值显存应不高于同 shape Ulysses。
  3. 性能报告必须同时给出 mode、backend、global/local sequence、head shape、warmup/repeat、
    forward、backward 和 peak allocated/reserved,不能只报告比例。

历史同 executor 的长序列结果可作为回归参照,但新 Trainer 的正式性能结论应由同环境同轮
A/B 重新生成,不直接沿用旧 Trainer 的数据。

9.5 工程检查

  1. git diff --check、Python compileall、pylint 和 Lizard 门禁通过。
  2. wheel 构建与隔离安装后,GDN module、CP executor、model adapter 和 Triton kernel 均存在。
  3. 不安装 Triton-Ascend 时可正常 import Hyper-Parallel 并运行 Eager。
  4. 选择 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,避免扩大当前
接口范围。

likedislike
fangwenyifangwenyi成员
29 天前 issue状态由 TODO 改变为 DONE
fangwenyifangwenyi成员
29 天前 关闭了 issue