【Task.32】FLA Chunkwise CP for GDN — RFC

[jambow0320]

1. 摘要

Relax 当前的 GDN（Gated Delta Net）上下文并行有两条路径：head 数能被 TP × CP 整除时走 MCore 原生 headwise CP（cp2hp all-to-all，每 rank 全序列 × 1/CP heads）；否则退回 _dcp_gdn_forward 的 all-gather fallback（运行时 wrapper 由 relax/backends/megatron/model.py:435-609 安装，路由条件见 model.py:491-510）——每个 GDN 层、每个 micro-batch 把完整投影序列 all-gather 到所有 CP rank（backward 对应 reduce-scatter），每 rank 重复执行整条卷积与 GDN scan，再切回本 rank 的 zig-zag 分片。fallback 语义正确，但通信量随全序列长度线性增长、GDN 计算与激活按 CP 倍数重复，并因此被 _assert_gdn_full_recompute()（model.py:339-365）强制 --recompute-granularity full；Qwen3.5 系列约每 4 层有 3 层为 GDN 层（参考 tests/utils/test_flops_counter.py 的 layer_types 配置），上述成本按层数放大。

本 RFC 提议把 NVIDIA Megatron-LM dev 分支已合入的 FLA chunkwise CP（KCP：每 rank 只持有并扫描连续时间分片，跨 rank 只交换卷积边界 token 与 (h_ext, M) 状态摘要）选择性回移到 Relax 固定的 MCore 基线上，并接入 Relax 的 per-microbatch 动态 CP、packed THD 元数据与既有回退路径。方案分两层：

1. FLA/MCore 兼容层（第 3 节）：FLA 0.4.1 → 0.4.2 精确升级；把上游 5139086e 中 Task 32 必需的 hunk 合入 docker/patch/ 的 MCore 补丁；不改变 Relax 任何默认路由，以 CP=1 / headwise / all_gather 三条同镜像回归门槛把"依赖升级风险"与"算法切换风险"隔离。
2. Relax 接入层（第 4 节）：暴露 --gdn-cp-mode={auto,chunkwise,headwise,all_gather}；auto 在 v1 对现有合法配置保持原路由，非法 headwise 几何改为提前报错；chunkwise 仅显式启用；落地动态 CP group 解析和 micro-batch 级元数据复用；回退路径见第 7 节。

本 RFC 重点覆盖：构造期能力声明与运行期路由分离；FLA 升级对现有 all_gather 路径本身数值的影响；权重同步/DCS 与 checkpoint 的 TP-only 分片不变式；以及"现有 8×H200 上真实 Qwen3.5-9B 的合法拓扑不会自然命中 fallback，需要使用 forced all-gather 受控对照"的实验设计（6.2）。

---

2. 上游 Chunkwise CP 能力与 Relax 现状

2.1 上游如何支持 Chunkwise CP

本 RFC 区分两个上游参考：

• Megatron-LM #3282 已合入 dev 分支，对应 commit 5139086e，提供 Chunkwise CP 的完整功能实现，是本任务的必选迁移基线；
• Megatron-LM #5664 仍为 Open，提出把 THD 布局路线提前构建并存入 PackedSeqParams，解决 5139086e 在每层重复计算路线的问题。本 RFC 将其中的路线预构建最小子集作为可选性能改进，不把整个未合入 PR 当成稳定依赖。

Chunkwise CP 的基础能力由 FLA 与 MCore 两层共同完成：

• FLA 0.4.2 提供 fla.ops.cp.build_cp_context，并让 causal_conv1d、chunk_gated_delta_rule 接受 cp_context。卷积只交换每个连续分片边界的 W-1 个 token；GDN forward all-gather (h_ext, M) 状态摘要并在本地折叠，backward 交换对应状态梯度。通信量由 head 几何决定，不再随完整序列激活线性增长。
• Megatron-LM 5139086e 负责把该能力接入 MCore：新增 context_parallel_layout.py 完成 zig-zag ↔ contiguous 转换；在 GatedDeltaNet.forward 中解析 CP group、构造 cp_context 并传给 FLA kernel；在 TransformerConfig 中增加 linear_cp_mode={headwise,chunkwise} 及对应的 head 整除校验。

上游执行链路为：

rank-local hidden_states
  → in_proj
  → zigzag_to_contiguous_chunks
  → build_cp_context(cu_seqlens, group, conv1d_kernel_size)
  → causal_conv1d(..., cp_context)
  → chunk_gated_delta_rule(..., cp_context)
  → contiguous_to_zigzag_chunks
  → out_proj

这条路径进入 GDN 前临时把 zig-zag 数据换成“每 rank 一段连续时间”，计算结束后再恢复 zig-zag，因此不会改变后续模型层的数据布局。上游同时支持 headwise 和 chunkwise：headwise 要求 heads 能被 TP×CP 整除；chunkwise 只要求能被 TP 整除。

5139086e 有两个需要 Relax 额外处理的点：

1. linear_cp_mode 是模型构造期的静态配置，而 Relax 的 CP group 和 CP size 可以随 micro-batch 变化，不能直接用一个静态字段决定所有 forward。
2. 每个 GDN 层的两次布局转换都会重复根据 cu_seqlens 计算相同的 token 搬运路线，并通过 .item() 等待 GPU 返回总 token 数；多个 GDN 层和 full recompute 会放大这部分开销。#5664 通过 prebuild_thd_cp_partition_routes 在数据准备阶段为每个 micro-batch 构建一次双向路线，后续层直接复用。

考虑到本任务目标包含长上下文性能，RFC 默认建议同时移植 #5664 的路线预构建最小子集。但该优化不影响 Chunkwise 的数值语义：若维护者希望只依赖已合入代码，可以不纳入这部分，保留 5139086e 的按需重建路线行为；其余 FLA/MCore 与 Relax 接入设计不变。

2.2 Relax 当前实现与补丁情况

依赖与本地补丁

Relax 镜像固定 Megatron-Bridge 2faedbf6 及其 MCore 85bced0a，安装 FLA 0.4.1；构建时把 Bridge 与 MCore 源码合并到 /root/Megatron-LM，再应用 docker/patch/latest/megatron.patch。当前 patch 对 GDN 有两处必须保留的兼容修复：

• Qwen3.6 的 _prepare_qkv_for_gated_delta_rule 会触发 torch.compile/Dynamo 失败，因此仅对该 helper 禁用编译并回到 eager；
• get_parameter_local_cp 使用 param[tuple(slices)] 表达多维基本切片，保证 headwise CP 的卷积参数、A_log、dt_bias 沿正确维度切分。

启动与模型构造

1. 参数校验阶段根据 rollout_max_context_len / max_tokens_per_gpu 推导本次任务可能使用的最大 CP，并向上取到 2 的幂；该值写入 args.context_parallel_size，作为静态最大 CP。
2. mpu.initialize_model_parallel(..., dynamic_context_parallel=True) 在初始化阶段创建静态 CP 组及可复用的动态 CP 子组；forward 内不会临时创建 process group。
3. 当前 MCore 没有 linear_cp_mode，TransformerConfig.__post_init__ 固定按 headwise 要求 heads 被 TP×静态CP 整除。Relax 在模型构造前安装 _relax_gdn_cp_config_assert：当 heads 只满足 TP 整除时，临时把 key/value head 数乘以 CP，使原生校验通过，随后恢复真实值。该操作只绕过构造检查。
4. 模型创建后，Relax 通过 _patch_gdn_for_dynamic_cp 替换 GatedDeltaNet.forward，安装运行时路由 wrapper。

每个 micro-batch 的执行流程

1. data.py 根据该 micro-batch 的最长样本和每 GPU token 预算计算 runtime_CP（1/2/4…），选择初始化阶段已创建的对应 CP 子组，并得到当前 rank 在组内的 cp_rank。
2. 数据随后进入两条 THD 准备路径之一：
   • 标准 pre-split THD：data.py 先按 cp_rank 对每条样本做 zig-zag 切分，再拼成 rank-local THD tensor，构造 batch["packed_seq_params"]（data.py:409-442）；
   • Bridge/VLM unsplit THD：Qwen3-VL 及共用该 Bridge 架构的文本 Qwen3.5/Qwen3.6 先接收未切分的 input_ids 和 attention mask，由 Bridge 在 embedding 后完成 CP/SP 切分，使用 batch["vlm_packed_seq_params"]（data.py:352-382、model.py:1039-1042）。
3. 标准路径把 cp_group、local_cp_size 和全局边界写入最终 PackedSeqParams；VLM 路径先把动态组和初始边界写入 vlm_packed_seq_params，Bridge repack 后的最终对象和边界仍需在接入层重新确认。
4. GDN wrapper 从该对象解析 runtime_CP/group/rank：CP=1 或非 THD 走原生路径；THD CP>1 且 heads 被 TP×runtime_CP 整除时调用 MCore headwise；否则进入 Relax all-gather fallback。
5. all-gather fallback 先在本地执行 in_proj，再收集所有 CP rank 的投影并按 zig-zag 规则重组完整序列；每个 rank 重复执行完整卷积和 GDN scan，最后切回本 rank 的 zig-zag 输出并执行 out_proj。该路径训练时要求 full recompute。

fallback 首次使用 GPU cu_seqlens 时会转成 CPU list，并以 _gdn_cu_seqlens_cpu 缓存在当前 PackedSeqParams（model.py:524-531）；同一 micro-batch 的后续 GDN 层直接复用，避免每层重复 .tolist() 和 D2H。该缓存只服务旧 fallback 的序列重组/切片，不包含 Chunkwise 的布局 route 或 cp_context。

因此上游与 Relax 的核心异同是：

维度	上游	Relax 当前	接入要求
FLA	0.4.2，提供 cp_context	0.4.1	精确升级到 0.4.2
CP 模式	静态 headwise / chunkwise	运行时 headwise / all_gather	保留现有默认路由，新增显式 chunkwise
动态 CP	resolve_cp_group 优先读取 PackedSeqParams.cp_group	data.py 每 micro-batch 选择 group/size，GDN wrapper 自行解析	收敛成单一解析逻辑，不在 forward 建组
zig-zag 与 contiguous THD转换路由	5139086e 每次交换重建；#5664（Open）提出在数据准备阶段预建	尚无 chunkwise 路由	默认纳入 #5664 的最小路线缓存子集；维护者不同意时可退回按需重建
数据布局	zigzag_to_contiguous_chunks 约定输入为 zig-zag	常规路径是 zig-zag；--allgather-cp 已是连续切分	拒绝 chunkwise 与 --allgather-cp 同开

2.3 迁移结论

5139086e 比 Relax 的 MCore 基线领先约 597 个提交，且 Relax 使用“Bridge + MCore 合并树 + 自有 patch”的装配方式，不能直接 cherry-pick。迁移应采用选择性 backport：

1. FLA 精确升级为 0.4.2，提供 build_cp_context 以及支持 cp_context 的卷积/GDN kernel；
2. 选择性回移 context_parallel_layout、GDN chunkwise forward、linear_cp_mode，并补齐旧基线缺少的 resolve_cp_group 等前置能力；
3. 合并上游 hunk 时保留第 2.2 节列出的两处 Relax GDN 兼容修复，避免 backport 覆盖现有可运行路径；
4. 为 GDN forward 增加单次调用级的 mode override，使同一个模型可按 micro-batch 选择模式而不改写共享 config；默认同时回移 #5664 的最小 THD 路线预构建接口，该项可按维护者意见独立移除；
5. 默认 auto 对现有合法配置保持原行为：CP=1/非 THD 走原生路径，THD CP>1 时可整除走 headwise、否则走 all_gather；非 THD 的非法 headwise 几何改为提前报错。chunkwise 仅显式启用，现有 all-gather fallback、TP-only 权重分片和 checkpoint 行为全部保留。

2.4 备选方案与取舍

• 整体升级 Megatron-Bridge/MCore：可直接获得较新的接口，但会同时改变 Attention、并行初始化、checkpoint、optimizer 和模型转换，难以隔离 Task 32 的收益与回归，因此不采用。
• 长期在 Relax monkey patch 中实现 Chunkwise：便于 PoC，但依赖 MCore 内部函数和参数布局，升级脆弱；生产 GDN conv/scan 与布局消费应回到 pinned MCore，Relax 只保留策略和 per-microbatch 元数据。
• 整体迁移 #5664：该 PR 尚未合入且覆盖 GPT/Hybrid/Attention/MTP/Qwen-VL，超出任务范围；只把 THD route 预构建最小子集作为可选 hunk。

下一节据此给出 FLA 升级、MCore patch 生成和逐文件 backport 方案。

---

3. FLA/MCore 兼容方案

3.1 FLA 升级与验证

Relax 当前固定 FLA 0.4.1，但该版本没有 Chunkwise CP 所需的 fla.ops.cp，卷积和 GDN kernel 也不接受 cp_context。因此 docker/Dockerfile:41 必须精确升级为 flash-linear-attention==0.4.2；

镜像构建后先做两类检查：

1. 能力检查：确认 build_cp_context 可以导入，并确认 causal_conv1d、chunk_gated_delta_rule 已支持 cp_context 参数；不满足则直接判定镜像构建失败。
2. 现有路径回归：FLA 升级不仅影响新增的 chunkwise，Relax 当前的 CP=1、headwise 和 all-gather 也在调用 FLA 的卷积/GDN kernel（model.py:560-598）。因此候选镜像必须分别与旧镜像比较这三条路径的 forward 和 backward；任何一条数值回归，都不能把问题归因于 chunkwise 接入，而应先阻断 FLA/MCore 兼容层合入（具体门槛见 3.4）。

报告记录最终安装的 FLA 版本、镜像 digest 和对应 tag commit ca910f88，避免后续重建环境时出现版本漂移。

3.2 Patch 方案

Relax 不直接在仓库中维护一份完整 MCore 源码。镜像构建时先组合固定版本的 Megatron-Bridge 2faedbf6 与 MCore 85bced0a，再把 docker/patch/latest/megatron.patch 应用到组合后的源码树（docker/Dockerfile:130-139）。因此 Chunkwise CP 的 MCore 改动也必须进入这份 patch。

上游 5139086e 比 Relax 的 MCore 基线新约 597 个提交，且 gated_delta_net.py、transformer_config.py 与 Relax 现有 patch 修改了相同区域，无法直接 cherry-pick。方案是在当前组合源码树上选择性移植第 3.3 节列出的能力，并重新生成一份包含“现有 Relax 修复 + Chunkwise CP backport”的完整 MCore patch。新的带日期 patch 保存到 docker/patch/megatron/，docker/patch/latest/megatron.patch 同步指向同一内容。

镜像构建必须验证 patch 能干净应用、无冲突标记，并能导入新增的布局模块和 GDN 模块。提交说明同时提供文件/hunk 级清单：说明哪些上游改动被采用、哪些被排除，以及排除原因，方便维护者确认没有夹带无关的 MCore 升级。

3.3 逐文件 backport 清单

这是 MCore patch 的核心范围：

目标文件	必须移植的能力	Relax 适配
megatron/core/context_parallel_layout	必选：5139086e 的 zig-zag ↔ contiguous 转换、THD rank 索引与 all-to-all 路由；可选但建议纳入：从 #5664 最小回移 route 的 build/get/prebuild/decode helper	采用可选优化时，布局转换优先读取预构建 route，未预构建时回退到 5139086e 的按需构建
megatron/core/ssm/gated_delta_net.py	解析运行时 CP group；进入/退出时转换布局；构造 cp_context；把它传给卷积与 GDN kernel	保留第 2.2 节两处 Relax 修复；增加 call-level mode override；仅在采用 #5664 可选优化时接入预构建 route；不移植 pre-GDR fusion 等后续优化
megatron/core/transformer/transformer_config.py	linear_cp_mode 字段，以及 headwise 按 TP×CP、chunkwise 按 TP 校验 heads 的规则	不采用上游默认 "chunkwise"；由 Relax 在构造期声明可用能力，实际模式仍由运行时分发器决定
megatron/core/packed_seq_params.py	必选：resolve_cp_group；可选但建议纳入：#5664 的 cp_partition_mode、zig-zag→contiguous / contiguous→zig-zag route tensor 字段	local_cp_size/cp_group 已存在于 85bced0a，无需重复 backport；维护者不接受未合入优化时，只补统一 group 解析
megatron/core/extensions/transformer_engine.py	修正混合模型中普通 Attention 层的 THD 输出长度：输入 tensor 含对齐 padding 时，TE 必须使用 padded cu_seqlens 分配同样长度的输出	当前 Relax 基线没有 pad_between_seqs 字段；只有本次 backport 同时采用该 padding 表示、且测试证明缺少此修复会导致输入/输出长度不一致时才纳入

上游通过 pyproject.toml/uv.lock 声明 FLA 依赖，但 Relax 的 MCore 是在 Docker 中组装而不是通过这些文件安装，因此同一依赖改动应落在 docker/Dockerfile，不复制上游的包管理文件。上游布局与 GDN 测试会改编到 tests/backends/megatron/；它们只验证被移植的基础能力，Relax 仍需第 5 节的自建测试覆盖动态 CP、VLM/unsplit 路径和现有 all-gather fallback。

#5664 改动范围很大，本任务不整体迁移。默认只采用与 THD 路由复用直接相关的最小子集：PackedSeqParams 的布局/route 字段，route 的构建与读取 helper，以及布局转换消费预构建 route 的接口；不迁移其 GPT/Hybrid 全局布局规划、Attention/MTP/Qwen-VL 改造等其他内容。兼容层负责提供这套 MCore 接口，Relax 接入层只负责在 data.py 构造 micro-batch 时调用它。由于该 PR 尚未合入，实施时固定并记录实际采用的 head commit；若维护者不接受未合入代码，则删除这组可选 hunk 和 data.py 预构建调用，继续使用 5139086e 的按需构建逻辑。

3.4 兼容层完成标准

本阶段只判断“FLA/MCore 能力是否正确迁移”，不验证 Relax 的动态 CP 路由和完整训练矩阵；后者统一放在第 5 节。完成标准为：

1. 新 patch 能干净应用，build_cp_context、布局转换模块和修改后的 GDN 可以正常导入；
2. CP=1 与静态 headwise 的 forward/backward 相对旧镜像无数值回归；现有 all-gather fallback 也需回归，因为它同样调用本次升级的 FLA kernel；
3. 最小 CP=2 Chunkwise 用例能够调用真实 FLA kernel，且 forward/backward 与 CP=1 参考结果满足容差，证明迁移的能力本身可用；
4. GDN 的 state_dict / sharded_state_dict 参数键与分片维度保持不变，证明迁移没有改变 TP-only 参数和 checkpoint 语义。

该阶段不修改默认 recipe，也不让 auto 选择 chunkwise。

---

4. Relax 接入方案

4.0 变更范围

文件	修改内容	目的
relax/utils/arguments.py	新增 --gdn-cp-mode={auto,chunkwise,headwise,all_gather} 及参数冲突校验，默认 auto	给用户显式选择和回退入口，同时保持旧命令行为不变
relax/backends/megatron/model_provider.py	仅对 GDN 模型在 provider.finalize() 前设置 MCore 的 linear_cp_mode 能力声明	让 TransformerConfig 用正确规则完成构造期 head 整除校验；普通 Attention 模型配置不变
relax/backends/megatron/model.py	把现有 GDN wrapper 重构为 dispatcher；在 Bridge/VLM 完成最终 repack 后验证实际 PackedSeqParams 的 padded 边界并准备 route	负责 per-microbatch 路由、call-level override，并保证 VLM 不会复用 embedding 前已经失效的边界/route
relax/backends/megatron/data.py	标准 THD 在 runtime CP 已确定后保留 CPU padded 边界，并在采用 #5664 时预构建 route；VLM/unsplit 只传递原始 host 长度，最终边界由模型入口在 repack 后确认	避免 GDN 层内新增 D2H，同时区分“可在 data.py 最终确定”的标准路径与“会被 Bridge 改写”的 VLM 路径
tests/backends/megatron/	增加第 5 节的布局、路由、真实 kernel、动态 CP、VLM 和 checkpoint 测试	同时验证上游能力与 Relax 特有接入，而不是只运行上游单测
scripts/training/sft/run-qwen3.5-9B-math-dynamic-cp-8xgpu.sh	暴露 GDN_CP_MODE、SEED、RUN_ID，并让不同实验使用独立日志/checkpoint 路径	作为真实训练入口运行四种模式和 before/after 对照，避免实验互相续跑或覆盖
scripts/benchmarks/	增加第 6 节所需的 step time、有效 tokens/s、峰值显存、GPU 利用率与通信时间采集	生成可复现的性能证据，而不是只报告单次日志中的平均速度

4.1 模式与参数

上游只提供模型级的静态 linear_cp_mode，但 Relax 的实际 CP size/group 会随 micro-batch 变化，而且还需要保留自有 all-gather fallback。一个共享静态字段无法同时表达这些运行时选择；直接在 forward 中改写它又会让交错 micro-batch 或 recompute 读到错误状态。因此本方案把“模型构造期采用哪套校验规则”和“本次 forward 实际走哪条路径”分开：前者只作为能力声明，后者由 Relax dispatcher 通过 call-level override 决定。

新增 --gdn-cp-mode {auto,chunkwise,headwise,all_gather}，默认 auto。四个模式的用途是：

模式	设计目的	运行行为
auto	保持现有合法 recipe 向后兼容	每个 micro-batch 继续使用当前规则：CP=1/非 THD 走原生路径；THD CP>1 且 heads 可被 TP×runtime_CP 整除时走 headwise，否则走 all-gather。不自动选择 chunkwise；非 THD 的非法 headwise 几何提前报错
chunkwise	显式试用和评测新路径	CP>1 时通过 call-level override 调用 FLA Chunkwise；不支持 deterministic 或 --allgather-cp，非法组合直接报错
headwise	回归对照和快速回退	强制调用 MCore headwise；当前 runtime CP 下 heads 不能整除时直接报错，不改走其他路径
all_gather	保留旧 fallback，作为不可整除几何的回退和正确性参照	CP>1 时仅支持 packed THD，强制调用现有 Relax all-gather 实现；训练继续要求 full recompute。CP=1 仍走普通非 CP 路径

构造期校验与运行时路由分离

MCore 回移 linear_cp_mode 后，以该字段替代当前 _relax_gdn_cp_config_assert 的临时 head 数改写。构造期只需选择能正确放行当前模型的校验规则：显式 headwise 或静态几何可整除的 auto 声明为 headwise；其余模式声明为 chunkwise，采用 TP-only 校验。后者只用于让 all-gather/Chunkwise 可用的不可整除模型完成构造。

实际执行路径仍由 dispatcher 按用户模式和每个 micro-batch 的 runtime_CP 通过 call-level override 决定；auto 继续保持“可整除走 headwise，否则走 all-gather”，不自动进入 Chunkwise。dispatcher 必须显式传入本次模式，不能依赖构造期默认值或在 forward 中改写共享 config；显式模式不做静默 fallback。

4.2 动态 CP 接入流程

动态 CP 要保证同一个 micro-batch 的数据切分、布局路线和所有 collective 使用同一 CP size/group。接入流程如下：

1. 初始化通信组：配置中的 context_parallel_size 表示最大 CP。初始化阶段创建最大组及 size 1/2/4… 的动态子组；forward 不创建新 group。
2. 选择 runtime CP：data.py 根据最长样本和每 GPU token 预算选择子组，把 cp_group 和 local_cp_size 写入 PackedSeqParams。二者必须同时存在且 local_cp_size == cp_group.size()。
3. 生成最终 padded 边界：每条序列必须满足 padded_length % (2×runtime_CP) == 0。标准 THD 在 data.py 中按 runtime CP padding/slice；VLM/unsplit 必须等 Bridge 在 embedding 后完成最终 repack，再使用最终 cu_seqlens_q_padded（没有时才用 cu_seqlens_q）检查，不能复用 repack 前的临时边界。
4. 准备元数据：保留最终边界的 CPU 形式并附到最终 PackedSeqParams，供 FLA build_cp_context 使用，避免每层 D2H。若采用 #5664，同时预构建双向 route；不采用时由 5139086e 在布局转换处按需构建。
5. 统一消费：dispatcher 通过 resolve_cp_group 从该对象解析 group/size/rank；MCore 布局转换和 FLA kernel 均使用同一结果。PackedSeqParams 随该 micro-batch 经过全部 GDN 层和 recompute，因此边界和可选 route 可复用。

自定义数据路径未提供 CPU 边界时，首个 GDN 层允许做一次转换并打印 warning。最终长度不满足 2×runtime_CP 时，在第一条布局 collective 前报错；动态 CP 正确性 gate 通过前不进行性能归因。

4.3 dispatcher 重构与旧路径收敛

_patch_gdn_for_dynamic_cp 重构为显式分发器：

def _gdn_forward_dispatch(self, hidden_states, attention_mask,
                          inference_context=None, packed_seq_params=None, *args, **kwargs):
    cp_size, cp_group, cp_rank = _resolve_and_validate_cp_group(self, packed_seq_params)
    requested_mode = get_args().gdn_cp_mode
    is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == "thd"

    # CP=1 不需要 CP 通信。
    if cp_size == 1:
        return _call_orig_with_dynamic_cp(...)

    if requested_mode == "auto":
        if not is_thd:
            mode = "headwise"  # 保持当前原生路径；非法 head 几何由下方校验拒绝
        else:
            mode = "headwise" if self.num_key_heads % (self.tp_size * cp_size) == 0 else "all_gather"
    else:
        mode = requested_mode

    if mode == "headwise":
        _validate_headwise_divisibility(self, cp_size)
        return _call_orig_with_dynamic_cp(..., linear_cp_mode_override="headwise")
    if mode == "chunkwise":
        # patched MCore 从 packed_seq_params 读取预构建路线
        return _call_orig_with_dynamic_cp(..., linear_cp_mode_override="chunkwise")
    return _all_gather_gdn_forward(...)  # 现 _dcp_gdn_forward 主体，原样保留

auto 解析出 headwise 后也必须显式传入 linear_cp_mode_override="headwise"。否则，当构造期为放宽校验而设置了 linear_cp_mode=chunkwise、但某个较小 runtime_CP 又满足 headwise 整除时，直接调用原始 forward 会误走 Chunkwise。各分支的实际执行路径由第 5.2 节测试观测验证。

现有代码按以下方式处理：

• all-gather fallback 的实现保持不变：_AllGatherFullSequence 负责带 autograd 的全序列收集，gdn_reassemble_full 恢复顺序，gdn_cp_slice 切回本 rank 输出；它继续作为显式 all_gather 模式和 Chunkwise 的 CP>1 正确性参照；
• _assert_gdn_full_recompute 只在 all-gather 分支执行，因为只有该路径在每个 CP rank 保留完整序列激活；headwise/Chunkwise 不应继承这项限制；
• 删除 _relax_gdn_cp_config_assert 的临时 head 数改写，由第 4.1 节的正式 linear_cp_mode 构造校验取代；
• 用统一的 _resolve_and_validate_cp_group 取代当前 _resolve_gdn_cp：内部调用 backport 的 resolve_cp_group 选择动态或静态组，再派生并校验 size/rank，避免 Relax 与 MCore 各自维护一套 group 解析逻辑。

---

5. 正确性验证计划

参照基准为 CP=1 全序列 GDN；CP>1 另外使用现有 all-gather 路径交叉核对。数值门槛沿用对应上游测试：

层级	比较对象	门槛
FLA GDN kernel	output、dq/dk/dv/dg/dbeta	归一化 RMS 误差 < 2e-3（FLA test_cp_gdn.py）
FLA causal conv	output、dx/dw/dbias	归一化 RMS 误差 < 1e-3（FLA test_cp_conv.py）
MCore GDN，CP=1	候选镜像 vs 旧镜像	atol=2e-4, rtol=2e-3，cosine similarity ≥ 0.99999
MCore GDN，CP>1	Chunkwise/headwise/all-gather vs CP=1	atol=2e-3, rtol=1e-2，cosine similarity ≥ 0.9999

所有用例同时要求有效 token 数和 mask 精确相等，且输出、梯度无 NaN/Inf。

上述阈值以 FLA 0.4.2 和 Megatron-LM 5139086e 的现有测试为起点；若目标 bf16 几何需要调整，必须在正式 before/after 运行前根据固定输入的误差分布一次性冻结并记录，不能在看到候选结果后放宽。

5.1 布局与元数据单元测试

• CP∈{1,2,4,8}，覆盖单样本、多条不等长样本、每样本 padding 和独立 padding 尾；每个 rank 的 zig-zag→contiguous→zig-zag 必须逐 token 往返恒等；
• 断言 MCore get_thd_context_parallel_rank_indices(..., "zigzag") 与 Relax slice_with_cp / gdn_reassemble_full 的 token 所有权和顺序一致，避免两边对 zig-zag 定义不同；
• 若采用 #5664：预构建 route 与按需 route 结果一致；同一 PackedSeqParams 跨层/recompute 复用，换 micro-batch 或 CP 子组后使用新 route；
• 动态 CP 的标准 THD 与 VLM/unsplit 最终 PackedSeqParams 均验证 diff(route_cu) % (2×runtime_CP) == 0；测试包含 embedding/repack 前后边界变化，确保 route 绑定最终对象；
• 覆盖可整除/不可整除几何的 TransformerConfig 构造，确认第 4.1 节的构造校验能够替代 _relax_gdn_cp_config_assert。

5.2 真实 kernel 分布式测试（CUDA/NCCL）

真实 CUDA/NCCL 测试分两层：

1. 直接调用带 cp_context 的 FLA causal_conv1d 和 chunk_gated_delta_rule，与无 CP kernel 比较输出和各输入梯度；序列长度覆盖 63/64/65 等 kernel chunk 边界，以及卷积宽度 W 附近的短分片。
2. 调用完整 MCore GatedDeltaNet，覆盖静态 SBHD（CP>1 时 batch=1）以及 packed THD 的单样本、多不等长样本、padding 尾；比较模块输入梯度和 conv/in_proj/out_proj/A_log/dt_bias 参数梯度。

几何覆盖 GQA（value heads 为 key heads 的整数倍）、key/value head dim 不同、TP/SP 开关和不可整除合成几何（如 4 key heads @ TP2/CP4）。四种用户模式必须通过 branch hook/counter 断言实际执行分支，不能只凭数值相近判断路由正确。另测试 chunkwise + --allgather-cp、deterministic、动态 SBHD、initial_state / output_final_state=True 和非法 head 几何均按设计报错。

5.3 并行 / 数据矩阵

TP	CP	SP	输入	说明
1	1	off	SBHD + THD	参照基准
1	2	off	SBHD(b=1)；THD 单/多样本	最小 CP
2	2	on/off	THD	TP×CP 交叉
1	4	off	THD	中等 CP
2	4	on	THD	主拓扑（单节点 8×H200）
1	8	off	THD	布局 / 长上下文压力

动态 CP：静态 CP=4 组内 cp∈{1,2,4} 混合 micro-batch，断言逐 mb 使用正确子组、无跨子组 collective；四种模式在该矩阵下的路由分布可复现。

5.4 checkpoint 与 recompute

• 同拓扑 save/resume：优化器与调度器状态、续跑首 step loss 连续；
• 权重 CP=1 保存 → CP>1 加载、CP>1 保存 → CP=1 加载；以 sharded_state_dict 键集合和分片维度对旧镜像逐项相等，验证 GDN 参数仍为 TP-only；
• 跨 CP 优化器状态恢复仅在现有分布式 checkpoint 已支持的范围内验证，不新造通用 resharder，超范围如实记录为限制；
• full recompute recipe 全程通过；chunkwise 下 recompute 开/关 parity（显存允许时）。

真实连续性使用同一 Qwen3.5-9B 配置做两条对照：A 连续训练 30 step；B 训练 20 step 后保存，在新进程恢复并继续 10 step。比较 step 21–30 的数据顺序、学习率、loss、grad norm 和最终参数；该用例通过后才满足摘要中的生产 gate。

5.5 端到端连续性

主 recipe 短训（≥30 step，CP=4）：auto 与显式 chunkwise 的 loss 曲线、grad norm、有效 token 数三线对照；有效 token 数逐 step 精确相等。

另选一个不含 GDN 的现有 CP recipe 做冒烟回归：默认 auto 下不得进入任何 GDN mode/route 逻辑或预构建 THD route，确认普通 Attention CP 路径与当前版本一致。

---

6. 性能与质量计划

6.1 冻结环境与运行纪律

• 模型/数据：Qwen3.5-9B 与 OpenMathReasoning-mini 使用固定 revision/文件 hash，具体值在首个基线运行前记录；
• 硬件：使用一台单节点 8×NVIDIA H200 完成全部正式对照；不跨节点、不混用 GPU SKU 或互联拓扑。首个基线运行前记录每卡显存容量、驱动/CUDA 版本、nvidia-smi topo -m、CPU/RAM 和镜像信息；
• 运行：每配置 220 step（前 20 warmup 剔除、后 200 计量）；paired seeds 1234/1235/1236，同一重复内所有模式用同 seed 同数据序；每配置×每重复独立 RUN_ID 命名空间；测量期禁 resume（resume 测试独立目录、不计入性能）；显存峰值统计 warmup 后 reset，启动期峰值单列；
• 每种模式使用独立 Ray job/新进程，避免 compile cache、CUDA allocator 和残留模型状态污染对照；三组 paired run 轮换执行顺序（正序、逆序、交错），降低机器温度和时间漂移带来的系统偏差；
• 全部失败/OOM 运行保留并报告；每次运行记录 Relax commit、镜像 digest、MCore patch digest、FLA 版本和 recipe。

6.2 实验设计

E1（主实验：四方对照）。使用真实 Qwen3.5-9B、TP2/CP4/SP on 和同一份 SFT recipe：

1. 旧镜像默认模式（当前该几何走 headwise），作为完整 before 基线；
2. 候选镜像 forced headwise，隔离 FLA 升级和 MCore forward 重构的依赖影响；
3. 候选镜像 forced all_gather；
4. 候选镜像 forced chunkwise。

比较 ④ vs ③ 衡量“用 Chunkwise 替换 fallback”的收益；④ vs ② 判断相对现有高效 headwise 是否非劣；② vs ① 单独量化依赖升级影响；④ vs ① 给出完整 before/after 结果。③ 与④首先使用相同 full-recompute 配置以保证公平，随后可增加④关闭 full recompute 的容量/性能实验。all-gather OOM 时作为容量结果报告，不虚构吞吐比。

Qwen3.5-9B 在现有 8×H200 资源的合法拓扑下会自然走 headwise，因此 forced all-gather 是受控算法对照，不代表默认线上路由。若后续能获得 32 卡，追加 TP2/CP16 的真实不可整除实验；否则用第 5 节不可整除合成几何完成正确性与 kernel 级性能验证，不把随机初始化模型作为 200-step 质量主实验。

E2（动态 CP）。先以真实模型连续 micro-batch 覆盖 cp∈{1,2,4}，要求最终 padded 边界全部满足 2×runtime_CP 整除且 forward/backward 通过；该正确性 gate 未通过时只报告 blocker，不进行动态 CP 性能归因。通过后再运行 dynamic-CP recipe，对比候选镜像 auto 与显式 chunkwise，并报告逐 micro-batch CP 大小分布。

E3（单层 shape crossover）。在同一 8×H200 TP2/CP4+SP 环境中，对单层 GDN 的 8K、16K、32K 序列分别比较 headwise 与 Chunkwise，使用独立进程和轮换顺序。该实验不替代完整模型结果，只用于识别布局转换与状态摘要通信随 shape 变化的性能拐点。

6.3 指标与统计口径

• 有效吞吐（主指标）：Σ(计量步有效非 padding token) / Σ(计量步各 rank 最大同步墙钟)；
• step time：mean / median / p95 / std；总 tokens/s 另列；
• GPU 利用率：nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used --format=csv,noheader,nounits --loop=1 1 Hz 逐卡原始样本留档；
• 显存：warmup 后 reset 的 torch.cuda.max_memory_allocated / max_memory_reserved，8 张 H200 取最大；
• 通信：抽样 step 的 profiler trace 中 layout all-to-all 与 (h_ext, M) all-gather 用时；若采用路线预构建优化，同时确认 GDN 层内没有 route metadata 引发的重复 D2H/host wait；
• 训练侧：逐 step loss、grad norm、精确有效 token 数。

6.4 判定阈值（首个基线运行前冻结；开跑后如需变更则全配置重跑）

• fallback 替换通过：④/③ 的有效吞吐比几何平均 ≥ 1.10 且逐对 > 1.00；或 ④/③ 的峰值 allocated 显存比 ≤ 0.90 且逐对 < 1.00、同时吞吐比 ≥ 0.95；
• 如果③ all-gather 在相同工作负载下 OOM 而④ Chunkwise 完成，则记为 fallback 容量 gate 通过，同时保留 OOM 证据，不计算吞吐比；
• headwise 非回归：④/② 吞吐比 ≥ 0.95，且④/②峰值 allocated 显存比 ≤ 1.05；该几何的 headwise 本就没有重复计算，不要求 Chunkwise 必须更快；
• 依赖非回归：②/①吞吐比 ≥ 0.95，且②/①峰值 allocated 显存比 ≤ 1.05；超出时作为 FLA/MCore 依赖变化单独分析；
• 有效 token 数逐配置精确相等；全部逐对比值公开发布，不只报均值。

6.5 质量非回归

• 训练 loss：计量段（step 21–220）按有效 token 加权平均，④相对①/②/③逐 seed 差值 ≤ 0.01 nats/token；对齐的逐 step 曲线与 grad norm 随报告发布；
• held-out 质量：从同源数据固定 512 行 held-out（训练集显式剔除并记录切分 hash），确保每个 220-step 运行保存最终 checkpoint，再离线评 token loss（Σ label_nll / Σ valid_label_tokens）；④相对①/②/③逐 seed 差值 ≤ 0.01 nats/token。

6.6 报告与证据

最终报告包含：全部固定修订（Relax commit、镜像 digest、patch digest、FLA 版本、模型/数据修订）、完整命令与并行拓扑、第 5 节全部正确性结果与误差分布、6.2 全部运行（含失败/OOM）、逐对比值与对齐曲线、已知限制与第 7 节回退命令。

---

7. 回退路径

1. 先回到旧执行路径：recipe 移除显式 --gdn-cp-mode chunkwise，恢复默认 auto。新版本的 auto 与当前 Relax 行为一致：可整除时走 headwise，否则走 all-gather。排查时也可显式指定 headwise；不可整除几何则指定 all_gather，并保留现有 full recompute 配置。
2. 再回到旧兼容环境：如果 CP=1/headwise 也出现回归，说明问题可能来自 FLA 0.4.2 或 MCore backport，而不只是 chunkwise 路由。此时恢复任务开始前记录的旧 Relax commit 和对应训练镜像（包含旧 MCore patch 与 FLA 0.4.1），不要混用未经验证的新源码与旧镜像。
3. 最后回滚代码：确认需要永久撤销时，先排空运行中的训练任务，再回滚 Relax 接入改动；若兼容层本身有问题，再一并回滚 MCore patch 与 FLA 升级。回退后运行 CP=1 和 headwise 冒烟测试。

正确性与性能报告记录新旧 Relax commit、镜像 digest、MCore patch digest、FLA 版本及对应 recipe，保证上述回退组合可以复现。

---

8. 风险与缓解

风险	缓解
FLA 0.4.2 升级改变现有三条路径的数值/性能	3.4 同镜像三路径回归门槛（含 all_gather）
选择性 backport 偏离上游或遗漏前置能力	固定上游 commit，提供 hunk 级 include/exclude 清单，移植对应上游测试
#5664 尚未合入且接口继续变化	只采用 route 预构建最小子集并固定实际 head commit；该组 hunk 可独立删除，退回 5139086e 的按需构建
构造期 linear_cp_mode 与运行时实际模式不一致	CP>1 时 dispatcher 始终显式传入 call-level override，并以 branch hook 测试四种模式
动态 CP group 解析不一致导致 collective 挂死	所有模块使用同一 resolve_cp_group 结果；group 只在初始化阶段创建；显式模式不静默回退
动态 CP 的最终 packed 长度不满足 2×runtime_CP，或 route 绑定了 VLM repack 前的旧边界	在最终 PackedSeqParams 上选择 padded 边界并做整除校验；真实 cp∈{1,2,4} gate 通过后才做动态性能实验
packed 边界或布局路线错误导致跨样本状态污染	多样本、padding 边界和 VLM/unsplit 路径均与 CP=1 比较输出及梯度
Chunkwise 的状态通信/布局转换开销抵消收益	进行第 6 节四方 200-step 对照；未达阈值则保持 auto 现状并使用 headwise/all-gather 回退

9. 请求维护者决策

1. 是否接受以 5139086e 为功能基线，选择性 backport 到 Relax 固定的 MCore 85bced0a，而不整体升级 Megatron-Bridge/MCore？
2. 是否接受第 4.1 节的构造期校验与运行时路由分离：linear_cp_mode 只用于模型构造，CP>1 的实际模式始终由 dispatcher 通过 call-level override 显式传入？
3. 是否接受 --gdn-cp-mode={auto,chunkwise,headwise,all_gather} 的命名和语义，其中 auto 保持当前 headwise/all-gather 路由，不自动启用 Chunkwise？
4. 现有资源是 8×H200，Qwen3.5-9B 在 TP2/CP4 的 auto 模式下会选择 headwise，不会选择 all-gather。计划在同一配置下分别强制运行 all_gather 和 chunkwise，用来比较 fallback 替换收益，同时保留 headwise 作为当前路径的对照。这个实验能否作为主要性能证据？如果要求使用 auto 模式本身就会选择 all-gather 的真实 Qwen3.5-9B 配置，则需要 32 卡 TP2/CP16；8 卡的不可整除合成配置只用于正确性和单层性能测试。
5. 是否接受默认回移尚未合入的 #5664 中 THD route 预构建最小子集？若不接受，将删除该组可选 hunk 与 data.py 调用，保留 5139086e 的按需路线构建。

[xiaoliang0601]

RFC 通过，可以进行后续开发。

1. 应该以 5139086e 为功能基线，选择性 backport 到 Relax 固定的 MCore 85bced0a，不整体升级 Megatron-Bridge/MCore。
2. 不接受。你的设计有些复杂，v1 版本先不做 per-call，v1 请采用启动级静态 --gdn-cp-mode：进程启动后固定 GDN CP 执行模式。linear_cp_mode 直接作为 MCore 构造期和运行期的配置使用，避免 Relax 侧 monkey patch，也避免 forward 期间修改 shared config/self 状态。dynamic CP 只允许动态改变 cp_group / local_cp_size，不动态改变 GDN CP mode。
3. 接受 --gdn-cp-mode={auto,chunkwise,headwise,all_gather} 的命名和语义；auto 在可走 headwise 时走 headwise，否则走 all_gather；不自动启用 chunkwise，但是 auto 是静态的、固定的模式，不要考虑 per-call。
4. 接受 8×H200、Qwen3.5-9B、TP2/CP4 下的 forced all_gather vs forced chunkwise 作为 fallback 替换收益的受控证据，同时保留 forced headwise 作为当前高效路径对照。但是，报告必须清楚说明：该拓扑下 auto 自然会走 headwise，因此 forced all_gather 不是默认线上路径，只是 fallback 替换实验。主要性能报告应采用四方对照：old image default、candidate forced headwise、candidate forced all_gather、candidate forced chunkwise。8 卡不可整除合成配置只用于正确性和单层机制验证，不作为真实 200-step 质量主实验。
5. 暂时不接受 #5664 作为 v1 版本的开发内容。v1 版本只需要依赖 #3282（5139086e），#5664 作为后续优化。

强调：

1. 正确性 gate 先于性能 gate。测试需要覆盖 CP=1、TP2/CP4、dynamic CP {1,2,4}、packed THD 多样本、padding tail、recompute、checkpoint save/resume。非法组合如 chunkwise + --allgather-cp、deterministic、unsupported inference 应该 fail-fast。
2. MCore patch 需要给出文件级和 hunk 级清单：哪些来自 5139086e，哪些是 Relax 适配，哪些现有 Relax GDN 修复被保留。不要夹带超出 v1 的依赖内容。
3. VLM / Bridge unsplit THD 路径要小心：chunkwise route / cp_context / padded cu_seqlens 必须绑定 embedding/repack 后的最终 PackedSeqParams，不能复用 repack 前的临时边界。
4. dynamic CP 只负责动态 cp_group / local_cp_size。所有 GDN、layout、loss/recompute 相关路径必须使用同 1 个最终 PackedSeqParams，并校验 local_cp_size == cp_group.size()；forward 过程中不得创建新的 process group。
5. auto 必须在启动 / 模型构造前一次性 resolve 成固定模式，并在日志里打印最终 resolved_gdn_cp_mode、context_parallel_size、tensor_model_parallel_size、linear_num_key_heads、linear_num_value_heads。训练过程中不要按 micro-batch 改变 GDN CP mode。

开发过程中有任何问题，都可以 Issue 中提出，我会关注。

祝顺利。

[jambow0320]

@xiaoliang0601 Hi~我已参考review comment，提交了第一阶段PR #251
主要用来升级FLA和迁移上游MCore能力，非常希望得到你的建议；该PR合入后，我会尽快开始第二阶段即Relax中chunkwisecp逻辑的开发工作；（如果你希望使用一个完整的pr而非这种分段式的pr方案，我也可以等开发完后提交一个总pr）；
谢谢！

[xiaoliang0601]

@xiaoliang0601你好~我已参考评论评论，提交了第一阶段PR #251 主要是为了升级FLA和迁移上游MCore能力，非常得到你的建议；该PR合入后，我会加速开始第二阶段即Relax中chunkwisecp逻辑的开发工作；（如果你希望使用一个完整的prf这种分段式的希望的pr方案，我也可以等开发完成后提交一个总pr）； 谢谢！

整体没什么问题，可以继续开发，comment 记得看下。

[jambow0320]

感谢review，我已经根据comment意见完善了PR，接下来我会基于第一个PR继续开发第二部分Relax侧的相关逻辑~

[jambow0320]

@xiaoliang0601 Hi，我在开发过程中碰到一个设计问题，想向你确认一下应该使用哪种方案呢？

gdn_cp_mode=auto 应该在哪一层 resolve

想请维护者确认 --gdn-cp-mode=auto 应该在哪一层解析成具体的 linear_cp_mode。

先说明旧实现：Relax 在模型构造前 monkey-patch TransformerConfig.__post_init__，临时调整 GDN head 数以绕过 heads % (TP×CP) 校验；模型构造后再 monkey-patch GatedDeltaNet.forward，由 forward 根据 runtime CP 动态选择原生 Headwise 或 Relax All-gather。

因此，旧版已经使用 MCore 构造阶段作为统一拦截点，但它没有在构造期 resolve mode，实际路由仍发生在 forward。

目前有两个可选方案。

方案一：让 MCore TransformerConfig 接受并 resolve auto

Relax 只负责把用户参数透传给 MCore：

linear_cp_mode = gdn_cp_mode

MCore 在 TransformerConfig.__post_init__ 的 GDN 校验前处理：

if (
    experimental_attention_variant == "gated_delta_net"
    and linear_cp_mode == "auto"
):
    divisor = tensor_model_parallel_size * context_parallel_size
    headwise_ok = (
        linear_num_key_heads % divisor == 0
        and linear_num_value_heads % divisor == 0
    )
    linear_cp_mode = "headwise" if headwise_ok else "all_gather"

完成构造后，MCore config 中只保留具体的：

headwise / chunkwise / all_gather

优点

1. 所有真正构造 MCore TransformerConfig 的路径使用相同规则：

   • Bridge：在 provider.finalize() 前传入 linear_cp_mode="auto"；
   • raw：core_transformer_config_from_args() 自动复制 args.linear_cp_mode；
   • custom：只要它构造 MCore config，就只需传入 requested linear_cp_mode，不需要自己读取 head 几何并复刻 resolve 逻辑。

2. 可以彻底删除 Relax 对 TransformerConfig.__post_init__ 的 monkey patch。

3. resolve 发生在完整模型几何已经可见、构造校验尚未执行的位置。

4. Relax 可以在三条路径汇合后的 setup_model_and_optimizer() 中读取最终的：

   module.config.linear_cp_mode

   并统一记录：

   requested_gdn_cp_mode=auto
   resolved_gdn_cp_mode=headwise/all_gather
   

缺点

1. MCore config 需要认识 "auto"，引入了一层上层路由策略，耦合更高。

2. MCore 构造会修改 linear_cp_mode，requested/resolved 的报告只能在模型构造后完成，或者另外保存 requested 字段。

3. 完全不使用 MCore config 的 custom model 不在该机制覆盖范围内。不过这种模型本来也不应该使用 MCore GDN CP mode。

方案二：MCore 只接受三个具体 mode，resolve 保留在 Relax

MCore 继续只接受：

headwise / chunkwise / all_gather

Relax 的 resolve_gdn_cp_mode() 在构造 MCore config 前返回具体值。

各路径需要分别接入。

Bridge

在 provider.finalize() 前 resolve：

provider.linear_cp_mode = resolve_gdn_cp_mode(args, provider)
provider.finalize()

raw/non-Bridge

在调用 core_transformer_config_from_args(args) 前，根据 args 中的 GDN 几何 resolve：

args.linear_cp_mode = resolve_gdn_cp_mode(
    requested=args.gdn_cp_mode,
    tp=args.tensor_model_parallel_size,
    max_cp=args.context_parallel_size,
    key_heads=args.linear_num_key_heads,
    value_heads=args.linear_num_value_heads,
)

config = core_transformer_config_from_args(args)

custom provider

如果 GDN 几何已经暴露在 args 中，可以和 raw 一样处理。

如果几何只在 custom provider 内部可见，就需要定义 provider 契约。Relax 将 resolve_gdn_cp_mode callback 传给 custom provider：

model = custom_model_provider(
    ...,
    gdn_cp_mode_resolver=resolver,
)

custom provider 必须在构造 TransformerConfig 前调用：

linear_cp_mode = gdn_cp_mode_resolver(
    tp=tp,
    max_cp=max_cp,
    key_heads=key_heads,
    value_heads=value_heads,
)

config = TransformerConfig(
    ...,
    linear_cp_mode=linear_cp_mode,
)

模型构造后，Relax 再检查实际 GDN config，防止 custom provider 静默忽略 mode。

优点

1. MCore config 保持独立，只认识实际执行算法。

2. requested/resolved 的责任都在 Relax，语义比较明确。

3. 不需要修改 MCore 对 linear_cp_mode 的公共接口。

缺点

1. Bridge、raw、custom 有不同的构造时机，需要在不同位置注入。

2. custom provider 必须新增并遵守构造前 resolver 契约。

3. resolve 逻辑更容易在不同 provider 路径中出现遗漏或行为不一致。

4. 当前如果只接入 Bridge，raw/custom GDN 的显式 mode 可能被忽略。

我的倾向

• 如果目标是静态 mode，并希望 Bridge/raw/custom 的 MCore 构造行为统一，方案一更简单，也比旧 __post_init__ monkey patch 更正式。
• 如果维护者认为 "auto" 属于 Relax policy、不应进入 MCore，则采用方案二，但需要明确 custom provider 的接口契约，不能只实现 Bridge 路径。

[xiaoliang0601]

我其实倾向于先不引入 auto，默认就走 headwise，不满足并行条件就 fail fast，让用户手动选择 chunkwise 或者 all_gather，把复杂性留到 v2 去做，v1 尽可能给个简洁的实现。你觉得呢？

[jambow0320]

我其实倾向于先不引入 auto，默认就走 headwise，不满足并行条件就 fail fast，让用户手动选择 chunkwise 或者 all_gather，把复杂性留到 v2 去做，v1 尽可能给个简洁的实现。你觉得呢？

好的呀，那就不引入auto，默认就是headwise，然后传参的时候直接把gdn_mode透传给linear_cp_mode，或者干脆在relax侧也统一把这个参数叫成linear_cp_mode，这样传进去mcore就能直接build，不满足并行条件就fail掉

[xiaoliang0601]

可以。

[jambow0320]

@xiaoliang0601 我更新了第二阶段的PR 254，完成了Relax侧的逻辑接入，以及一些正确性和性能实验，希望得到review和建议~

[jambow0320]

更新了接入5664性能优化的代码，#273 chunkwise的cp性能要明显好于headwise和allgather

[xiaoliang0601]

@jambow0320 Hi，我们这边升级了 megatron，导致这个 RFC 有些重复的工作不能合入（集中在第一个 PR），其他 2 个 PR 是可以合入的，你这边还有推动合入的意愿吗？如有，我帮你推动，内部测试一下。

[jambow0320]

@xiaoliang0601 嗯嗯哈喽，有什么我这边能做的不，想持续推进合入的~

[xiaoliang0601]

我的打算是，我先把你的代码整理下，拉回内部跑通验证（基于内部的最新代码），然后你根据我整理的代码，提交个新的 PR，帮助我跑更广泛的测试，之后就可以合入了。

[jambow0320]

嗯嗯好的，有需要的话我这边随时能帮忙验证和推进，这个issue下ping我就行~