[Task32] feat: fla mcore chunkwise cp stage 1

[jambow0320]

Task 32 第一阶段：FLA 0.4.2 与 MCore GDN Chunkwise CP 兼容层

对应 RFC：redai-infra/Relax#213。

上游参考：NVIDIA/Megatron-LM#3282，merge commit 5139086e。

1. Summary

这个 PR 只完成第一阶段的依赖和 MCore 能力准备：

1. 将 flash-linear-attention 从 0.4.1 精确升级到 0.4.2；
2. 将上游 GDN chunkwise context parallel 能力选择性 backport 到 Relax 固定的 MCore 85bced0a；
3. 让 MCore TransformerConfig.linear_cp_mode 能正式声明三种构造模式：
   • headwise
   • chunkwise
   • all_gather
4. 保持现有 Relax 默认执行路径不变。

这个 PR 不修改 relax/ 目录中的运行时路由，也不新增 CLI 参数。--gdn-cp-mode、auto 解析和 all-gather dispatcher 收敛留给第二阶段。

为什么 MCore 需要声明 all_gather

MCore 自己只实现 headwise 和 chunkwise；all-gather 算法仍由 Relax 现有 fallback 实现。

但 all-gather 和 chunkwise 都只要求：

heads % TP == 0

而 headwise 要求：

heads % (TP × CP) == 0

如果 MCore 只有 headwise/chunkwise 两个声明值，Relax 为了构造 all-gather 模型只能：

• 假装配置是 chunkwise，实际运行 all-gather；或
• 继续 monkey-patch TransformerConfig.__post_init__，临时修改 head 数骗过校验。

两种方式都会让“构造时声明的算法”和“运行时执行的算法”不一致。

因此本 PR 增加 Relax-only all_gather 声明：

• 构造时使用 TP-only head 校验；
• CP>1 时如果误入 MCore 原始 forward，直接报错，说明 Relax wrapper 没有正确截获；
• 第二阶段可以让 auto 在模型构造前显式解析为 headwise 或 all_gather，不再 monkey-patch head 数。

2. Background

Relax 当前 GDN CP 有两条路径：

Headwise

通过 all-to-all 将“每卡一段序列”转换成“每卡一部分 head、完整序列”。

优点：

• 每个 token 只计算一次；
• 是当前 MCore 原生路径。

限制：

key/value heads 必须能被 TP × CP 整除

Relax all-gather fallback

当 head 无法被 TP×CP 整除时，Relax：

1. 在每张卡上执行 in_proj；
2. all-gather 完整序列；
3. 每个 CP rank 重复执行完整 conv 和 GDN scan；
4. 切回当前 rank 的 zigzag 分片；
5. 执行 out_proj。

优点：

• 只要求 head 能被 TP 整除；
• 能覆盖高 CP、少 head 的模型。

代价：

• 完整序列 activation 被复制到每个 CP rank；
• conv/GDN scan 重复 CP 次；
• 当前训练要求 full recompute。

Chunkwise

每个 rank 保留一段连续时间片和全部 TP-local heads，只跨 rank 交换：

• causal conv 所需的边界 token；
• GDN 分片末尾的状态摘要。

它同样只要求：

heads % TP == 0

但不会 all-gather 完整序列，也不会重复完整 scan。

3. Scope

本 PR 修改：

• docker/Dockerfile
• docker/patch/megatron/20260805-85bced0ae.patch
• docker/patch/latest/megatron.patch
• tests/backends/megatron/test_gdn_chunkwise_cp_layout.py
• tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py

4. Dependency change

FLA 0.4.1 → 0.4.2

FLA 0.4.2 是第一个提供以下接口的发布版本：

from fla.ops.cp import build_cp_context

并让：

causal_conv1d
chunk_gated_delta_rule

接受：

cp_context=

Relax 不通过 MCore 的 pyproject.toml/uv.lock 安装依赖，因此没有照搬上游这两个文件，而是在 Dockerfile 精确固定：

flash-linear-attention==0.4.2

FLA/MCore 能力由专项测试直接验证，不在 Dockerfile 做 import 检测：

• kernel 测试会实际导入并调用 build_cp_context；
• causal_conv1d 和 chunk_gated_delta_rule 会带真实 cp_context 执行 forward/backward；
• MCore 模块测试会实际走 layout、context build、deterministic reference 和 mode routing。

5. MCore backport

5.1 context_parallel_layout.py

这个文件从 5139086e 完整迁入，内容与上游一致，提供：

• THD/SBHD 的 zigzag → contiguous；
• contiguous → zigzag；
• packed token route；
• 实际 all-to-all 和 receive reordering。

5.2 packed_seq_params.py

上游 #3282 所在 dev 基线已经有 resolve_cp_group()，因此该文件不在官方 PR diff 中。

Relax 固定的 85bced0a 已经有：

PackedSeqParams.cp_group
PackedSeqParams.local_cp_size

但缺少统一解析 helper，所以本 PR额外 backport：

resolve_cp_group(static_cp_group, packed_seq_params)

行为：

• packed_seq_params.cp_group 存在时使用当前 micro-batch 动态组；
• 否则使用模型构造时的静态组。

这里的“动态组”不是在 forward 中临时创建。Relax 启动分布式环境时通过：

initialize_model_parallel(..., dynamic_context_parallel=True)

一次性创建最大 CP 组及允许的 size 1/2/4/... 子组。data.py 每个 micro-batch 只根据长度选择已经存在的子组：

cp_group = get_dynamic_data_context_parallel_groups(group_size=runtime_cp_size)

并把这个现成的 ProcessGroup 对象写进最终 PackedSeqParams。GDN forward 只读取并使用它，不调用 new_group，因此不会在训练过程中反复建组。

MCore 直接从最终 group 派生 size/rank，同时按 maintainer 要求，每次 GDN forward 都校验对象的：

local_cp_size == cp_group.size()

5.3 TransformerConfig.linear_cp_mode

从上游迁入：

• linear_cp_mode 字段；
• headwise 使用 TP×CP head 校验；
• chunkwise 使用 TP-only head 校验；
• linear-attention CP 与普通 Attention cp_comm_type 分离。

Relax 修改一：默认值保持 headwise

上游默认：

chunkwise

本 PR默认：

headwise

原因：只升级镜像不能让现有 recipe 静默切换算法。

Relax 修改二：增加 all_gather

合法值变成：

headwise / chunkwise / all_gather

all_gather 构造时使用 TP-only head 校验，但 MCore 不做该all_gather fallback实现，而是会在relax端劫持调用已经写好的实现。

因此 CP>1 时如果 MCore 原始 forward 收到 all_gather，会 fail-fast，防止漏装 Relax wrapper 后静默执行错误算法。

Relax 修改三：拒绝未解析的 auto

auto 只属于未来 Relax CLI，必须在构造模型前解析。MCore 构造时只接受三个具体值。

5.4 gated_delta_net.py

从上游迁入的核心能力：

• build_cp_context；
• headwise/chunkwise group 分流；
• seq_len_post_headwise / seq_len_global；
• zigzag ↔ contiguous；
• 将 cp_context 传给 conv 和 GDN rule；
• SBHD chunkwise batch 限制；
• non-packed CP context cache；
• chunkwise 每序列长度整除校验。

Relax 要求：最终 dynamic CP 元数据只校验一次

GDN 收到最终 PackedSeqParams 后检查：

• cp_group 与 local_cp_size 必须同时存在或同时为空；
• dynamic CP 下 local_cp_size 必须等于 cp_group.size()。

每次 GDN forward 都直接检查，不缓存“已验证”状态。这样逻辑更直观，也能覆盖 Bridge repack 后替换的新对象。

以下差异需要重点说明。

旧 MCore 兼容：保留单体 forward

官方 PR基于更新的 dev，GDN 已经拆成：

forward
  → _forward_compute
  → pre_gated_delta_rule
  → _a2a_cp_to_hp / _a2a_hp_to_cp

Relax 固定的 85bced0a 仍是单体 forward。

本 PR没有回移约 597 个提交间的整套 GDN 重构，而是在旧 forward 内嵌入最小 chunkwise 能力：

• inline 解析 mode/group；
• 用 cp_size_headwise 替换旧 self.cp_size；
• chunkwise 下 headwise 操作退化为 size 1/no-op；
• 在 conv/scan 前后执行布局转换。

这样避免同时改变旧 headwise 路径的整体函数结构。

Relax 兼容：_prepare_qkv_for_gated_delta_rule 参数可选

上游把 cp_size_headwise 设为必填参数。

现有 Relax all-gather fallback 直接调用这个 MCore 私有 helper，由于我们第一个pr不修改Relax中的逻辑，为了fallback调用时不出错:

因此本 PR暂时允许：

cp_size_headwise=None

未传时回退到 self.cp_size。现有 fallback 调用前会临时将 self.cp_size=1，含义是 all-gather 后不再沿 CP 切 head。

第二阶段修改 Relax wrapper 显式传 cp_size_headwise=1 后，可以再将该参数恢复成上游的必填形式。

旧 MCore 兼容：保留 helper 名称

上游将：

get_parameter_local_cp
→ get_parameter_local_cp_headwise

本 PR保留旧名称，只让 cp_group=None 表示 size 1，减少旧基线调用面变化。

与上游对齐：统一传 cp_context

本 PR给 FLA rule 和 torch reference rule(主要是deterministic模式使用) 都统一传：

cp_context=chunkwise_cp_context

非 chunkwise 时值为 None。

同时给 torch_chunk_gated_delta_rule 增加上游相同的 cp_context=None 参数，并断言它只能为 None。真正的 chunkwise CP 仍然只支持 FLA。

保留既有 Relax 修复

当前 Docker 每次只应用一份累计 MCore patch。新 patch 必须继续包含旧 patch 的有效修复：

1. _prepare_qkv_for_gated_delta_rule 周围禁用 Dynamo compile；
2. param[tuple(slices)] 的多维基本切片修复。

否则 latest 切到新 patch 后，旧修复会从镜像中消失。

6. Explicit exclusions from upstream PR

gdn_conv_pad_alignment

没有迁入。

官方 PR所在的新 dev 已有固定 conv alignment padding，并将其改为可配置；Relax pinned 基线没有这套优化。它不是 chunkwise 必需能力，而且 local chunk padding 与 chunkwise recurrence 不兼容。

transformer_engine.py padded boundary fix

没有迁入。

该修复解决的是通用问题：物理 THD tensor 比普通 cu_seqlens 描述得更长时，TE 可能创建过短输出。

当前 Relax 标准路径会把尾部 padding 直接追加到普通边界；unsplit handoff 也将普通和 padded 边界设为同一值，因此现有表示不触发该Bug。

第二阶段PR会在实际场景进一步测试该Bug。

新版 GDN 重构和融合

没有迁入：

• _forward_compute 整体拆分；
• selective recompute_gdn wrapper；
• streamed/fused pre-GDR；
  -新版 headwise A2A wrapper；
• helper rename。

原因是这些不是 chunkwise 最小依赖，且会显著扩大旧路径回归面。

#5664 route prebuild

没有迁入未合入的 NVIDIA/Megatron-LM#5664：

• 无 route tensor 字段；
• 无 prebuild_thd_cp_partition_routes；
• THD layout route 每次调用按 cu_seqlens 重建。

这是性能项，不影响数值语义。是否需要 route cache 留给实际 profiler 数据决定。

7. Tests

7.1 test_gdn_chunkwise_cp_layout.py

50 个 CPU item，来源和目的：

• 上游布局定义：所有 token 恰好覆盖一次；
• 上游边界检查：uneven、duplicate、decreasing、unknown layout；
• Relax 特有：MCore zigzag 与 slice_with_cp / gdn_cp_slice 逐 token 一致；
• 上游 config：headwise TP×CP、chunkwise TP-only；
• Relax config：默认 headwise、all_gather TP-only、拒绝未解析 auto；
• 静态模式约束：GatedDeltaNet.forward 不允许 per-call mode override；
• 旧基线前置能力：resolve_cp_group 优先动态组。

结果：

50 passed

7.2 test_gdn_chunkwise_cp_gpu.py

10 个 pytest item：

1. FLA kernel fp32：
   • conv output/dx/dweight/dbias；
   • GDN output/dq/dk/dv/dg/dbeta。
2. FLA kernel bf16：同一组生产 dtype 检查。
3. 完整 GDN fp32：CP1 vs headwise CP2 vs chunkwise CP2。
4. 完整 GDN bf16：同一完整模块生产 dtype 检查。
5. deterministic headwise：torch reference 接受 cp_context=None 并与 CP1 对齐。
6. recompute parity：外层 activation checkpoint 重放 chunkwise forward，输出和梯度不变。
7. TP2/CP2：4 GPU 同时覆盖 TP head 分片和两种 CP 模式。
8. 真实 layout round trip：THD/SBHD 的 NCCL zigzag → contiguous → zigzag 必须 bit-exact。
9. fail-fast：all_gather 漏 wrapper、dynamic CP group/size 不一致、chunkwise deterministic、inference 等非法组合。
10. checkpoint：state/sharded-state 分片不随 CP mode 变化，并实际执行 state_dict save/load。

运行结果：

• 首次全量：9 passed，deterministic 新用例因门槛 1e-5 过严失败；
• 实测 deterministic RMS 2.813e-5、cosine 1.0；
• 改用完整模块统一 fp32 门槛 1e-3 后，该项单独复跑通过。

即 10 个 item 均分别通过。

7.3 与上游测试的差异

比上游增加：

• MCore layout 与 Relax 数据路径逐 token 一致；
• 真实 NCCL layout round trip；
• FLA dweight/dbias/dg/dbeta；
• Relax-only all_gather 构造与 fail-fast；
• 新旧 CP 算法同权重对照；
• state_dict 实际序列化/恢复。

从上游补回：

• duplicate/decreasing boundary；
• deterministic torch reference；
• recompute parity；
• TP2/CP2。

仍留到第二阶段pr的测试——真实 recipe：

• dynamic CP {1,2,4} 混合；
• TP2/CP4 目标拓扑；
• Bridge 最终 repack 后 TE 输出长度；
• 完整 distributed checkpoint save/resume 首步 loss 连续性。

8. Controlled old/new numerical evidence

测试升级FLA版本是否会带来精度问题，相关代码没有放进本PR，在本地环境实验后仅将结论表述如下:

实验环境：

• 旧镜像：relaxrl/relax:latest，digest 前缀 sha256:2d8fce08a40e；
• 旧环境：FLA 0.4.1 + 20260506-85bced0ae.patch；
• 候选环境：同一基础镜像 + FLA 0.4.2 + Task 32 MCore patch；
• 参数按参数名固定 seed 重写；
• 输入和 backward seed 固定；
• 比较 forward、输入梯度和所有参数梯度。

结果：

• CP1：9 个 tensor，6 个 bitwise identical，0 个超容差；
• headwise CP2：18 个 tensor，17 个 bitwise identical，0 个超容差；
• Relax all-gather CP2：18 个 tensor，18 个 bitwise identical，0 个超容差；
• 合计：45 个 tensor，41 个 bitwise identical，0 个超容差；
• 最大 normalized RMS 约 5.7e-7；
• cosine 全部为 1.0。

非 bitwise 的差异只出现在少量 backward reduction，量级符合 FLA 版本改变归约顺序。

9. How Relax will integrate in phase 2

9.1 新增启动参数

--gdn-cp-mode={auto,headwise,chunkwise,all_gather}

9.2 启动时一次性解析 auto

使用静态最大 CP：

heads % (TP × max_CP) == 0
    → headwise

否则
    → all_gather

auto 暂不自动选择 chunkwise；chunkwise 只允许显式启用。

解析结果必须在模型构造前写入：

TransformerConfig.linear_cp_mode

并打印：

requested_gdn_cp_mode
resolved_gdn_cp_mode
TP
max_CP
key/value head 数

9.3 删除构造期 monkey patch

删除 _relax_gdn_cp_config_assert() 临时放大 head 数的方式。

linear_cp_mode="all_gather" 已能用正式 TP-only 规则完成构造。

9.4 收敛 runtime wrapper

现有 _patch_gdn_for_dynamic_cp 收敛为薄 dispatcher：

runtime CP=1
    → MCore 普通 forward

mode=headwise/chunkwise
    → MCore forward

mode=all_gather
    → Relax 现有 fallback

wrapper 不再修改：

self.cp_size
self.pg_collection.cp
self.config.linear_cp_mode

动态 group 只通过最终 PackedSeqParams.cp_group 传入。

9.5 all-gather fallback 显式适配新 helper

将现有调用改成：

_prepare_qkv_for_gated_delta_rule(..., cp_size_headwise=1)

随后可将 MCore 中该参数恢复成与上游一致的必填参数。

9.6 数据与 Bridge

• 标准 THD：data 层选择 runtime CP group，并保证每条 padded 序列满足 2×runtime_CP 整除；
• VLM/unsplit：必须等待 Bridge 最终 repack 后，再使用最终 PackedSeqParams 构造 cp_context；
• loss/recompute 不再单独依赖 batch["dynamic_cp_size"] / dynamic_cp_rank，而是从同一个最终 PackedSeqParams 派生 group/size/rank；
• forward 内不得创建 process group；
• 需要验证最终普通边界覆盖物理 THD 行数，决定是否迁入 TE padded-boundary fix。

9.7 第二阶段正确性 gate

在性能实验前完成：

• dynamic CP {1,2,4} 同一训练任务混合；
• TP2/CP4 目标拓扑；
• headwise/chunkwise/all_gather 实际分支计数；
• recompute 开关；
• distributed checkpoint save/resume；
• Bridge/VLM 最终 repack；
• chunkwise 与 --allgather-cp 等冲突参数 fail-fast。

10. Maintainer feedback requested

希望 maintainer 明确确认以下设计点：

1. 是否接受 linear_cp_mode="all_gather" 作为 Relax-only 构造声明：
   • MCore 不实现算法，而是保持现状，在Relax端调用fall_back逻辑；
   • CP>1 误入 MCore 时 fail-fast；
   • 用于第二阶段删除构造期 monkey patch。
2. Relax 默认保持 headwise，而不是采用上游默认 chunkwise。
3. 第一阶段为兼容现有 Relax fallback，将 cp_size_headwise 暂时保留为可选参数；是否接受第二阶段显式传 1 后再恢复上游必填签名。
4. 是否希望现在就迁入 transformer_engine.py 的 padded-boundary fix，还是等 Bridge 最终 repack 测试证明存在长度不一致后再迁。
5. 是否接受暂不迁入 gdn_conv_pad_alignment、新版 GDN 重构和 #5664 route cache。
6. FLA版本升级精度测试的双镜像数值结果是否只保留在 PR 文档，而不提交 CI 无法执行的 runner。

11. Test plan

如下命令已在本地8xH200运行并全部通过；

# CPU
pytest tests/backends/megatron/test_gdn_chunkwise_cp_layout.py -q

# GPU：其中 TP2/CP2 需要 4 张 GPU，其余使用 2 张
pytest tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py -q

# Repository checks
pre-commit run --all-files

[xiaoliang0601]

这里是如果是 headwise 会有问题吗？是否需要改成 cp_size=self.cp_size？

[jambow0320]

确实，headwise 后续也会执行 cu_seqlens_q // cp_size_headwise，原先传 cp_size_chunkwise 会在 headwise 下变成 1，导致跳过 per-sequence divisibility check (虽然其实在relax的逻辑中，headwise的cu_seqlens_q数据会在data.py中就提前check，所以其实不会出问题，但是为了统一和上游pr的异同，我们还是在这里把这块check加上)。已提交最新commit，改为传当前 resolved cp_size = cp_group.size()。
没有使用 self.cp_size，因为 dynamic CP 下它指的是静态最大 CP，可能大于当前 micro-batch 的 runtime CP，使用cp_size才是真正的runtime CP。