[Task32] feat: fla mcore chunkwise cp stage 2 by jambow0320 · Pull Request #254 · redai-studio/Relax

jambow0320 · 2026-08-10T10:23:35Z

Task 32 第二阶段：Relax 侧 GDN CP 路由接入

对应 RFC：redai-infra/Relax#213。以及一些实验结果参考from redai-infra/Relax#206

对应第一阶段 PR：redai-infra/Relax#251。（本PR基于第一阶段PR的修改上新增改动，即commit 22b1206）

根据 RFC issue 中的讨论结果，本PR保留最简单的启动级静态模式：

不接受 per-call 模式切换，也不引入 auto。Relax 直接暴露 MCore 同名的 --linear-cp-mode={headwise,chunkwise,all_gather}，默认 headwise；不满足 Headwise 几何时构造期 fail fast，由用户显式选择 chunkwise 或 all_gather。
v1 只依赖 MCore 已合入的 #3282（5139086e），暂不接受 #5664（THD route 预构建）。

1. Summary

本 PR 完成 Relax 接入层：

直接复用 Stage 1 MCore TransformerConfig 已暴露的 --linear-cp-mode={headwise,chunkwise,all_gather}，默认 headwise；Relax 不重复注册同名参数；
在 relax/backends/megatron/arguments.py 新增 _validate_linear_cp_mode，对 Chunkwise 新路径的非法组合提前报错；
raw/标准 custom 由 core_transformer_config_from_args() 自动读取 args.linear_cp_mode；Bridge 通过统一 config override 列表在 provider.finalize() 前透传同一字段；CP 启用后由 rank 0 打印最终实际 mode；
删除 relax/backends/megatron/model.py 中构造期的 _relax_gdn_cp_config_assert head 数 monkey patch；
保留原 _dcp_gdn_forward 和 All-gather 主体的位置，只把顶部路由改成读取静态 linear_cp_mode：CP=1、headwise、chunkwise 调用 patched MCore 原始 forward，all_gather 且 CP>1 继续执行原 fallback；
all-gather fallback 显式传 cp_size_headwise=1，不再临时改写 self.cp_size；
CPU 单测覆盖 CLI、参数校验、长度整除 fail-fast 和 dispatcher；真实端到端训练结果见 §12。

本 PR 不再修改 Stage 1 的 MCore config 语义：MCore 仍只接受三个具体 mode；不实现 auto，不实现 #5664 route 预构建。正确性与性能验证结果分别见 §11 和 §12。

2. Background

Stage 1 PR #251 只做了两件事：把 FLA 升级到 0.4.2；把上游 linear_cp_mode（headwise/chunkwise）能力和 Relax-only 的 all_gather 声明 backport 到 pinned MCore。Stage 1 没有改 relax/ 目录，训练时实际走的路由仍是 Stage 0 的运行时判定：

构造期：_relax_gdn_cp_config_assert 在 CP 生效时临时把 linear_num_{key,value}_heads 乘以 CP，骗过 MCore 的 heads % (tp*cp) 校验，构造完再还原；
运行期：_patch_gdn_for_dynamic_cp 装的 _dcp_gdn_forward 每次 forward 现算 native_ok = num_key_heads % (tp_size * cp_size) == 0，native_ok 时临时把 self.cp_size/self.pg_collection.cp 换成当前 micro-batch 的动态值再调用原始 forward，用完还原；native_ok 为假且是 packed THD 时才用 Relax 自己实现的 all-gather fallback。

即原本relax实现中"构造时声明的算法"和"运行时执行的算法"是分离的、且靠运行时反复读写共享的 self.cp_size/self.pg_collection.cp 来"借用"MCore 的原生 forward——第二阶段的任务就是把这套隐式判定换成第 1 节列的显式静态路由。

3. Scope

本 PR 修改：

relax/backends/megatron/arguments.py（新增 _validate_linear_cp_mode）
relax/backends/megatron/model_provider.py（Bridge override 列表透传 linear_cp_mode）
relax/backends/megatron/model.py（删除构造期 monkey patch；最小改动收敛 GDN wrapper；打印实际启用 mode）
tests/backends/megatron/test_gdn_chunkwise_cp_layout.py（更新具体 mode/CLI 命名相关断言）
tests/backends/megatron/test_gdn_cp_mode_stage2.py（新增，CPU）

4. `--linear-cp-mode` CLI 参数

Stage 1 的 MCore patch 已把 linear_cp_mode 加入 TransformerConfig；Megatron 参数系统会自动把该 dataclass 字段暴露为 --linear-cp-mode，默认值就是 headwise。Relax 如果再次 parser.add_argument("--linear-cp-mode", ...) 会产生 argparse option conflict，因此 Stage 2 只消费和校验现有字段，不重复定义。

默认 headwise。如果 key/value heads 不能被 TP*max_CP 整除，MCore 构造期直接报错；用户必须显式改用 chunkwise 或 all_gather。

5. 参数校验

文件：relax/backends/megatron/arguments.py，新增 _validate_linear_cp_mode(args)，在 validate_args 末尾调用。

规则	触发条件	原因
拒绝 `chunkwise + --allgather-cp`	`linear_cp_mode == "chunkwise" and allgather_cp`	chunkwise 需要 MCore 的 zig-zag THD 布局；`--allgather-cp` 把数据路径换成单一连续分片，两者结构冲突
拒绝 `chunkwise + --deterministic-mode`（当 CP 可能 > 1）	`linear_cp_mode == "chunkwise" and (dynamic_context_parallel or context_parallel_size > 1) and deterministic_mode`	chunkwise 的 deterministic 参考实现没有 CP context 路径

这里不新增 all_gather + CP>1 的 recompute 启动参数限制。旧版 _assert_gdn_full_recompute() 仍只在训练实际进入 all-gather fallback 时检查，并保留其 no_grad 跳过语义；把同一策略提前到模型无关的参数校验阶段会扩大旧行为边界，不是 Chunkwise 接入所必需的改动。

不在参数校验阶段实现：

Headwise 几何由 MCore TransformerConfig.__post_init__ 的既有 gate 校验；Relax 不复制 head 数判断。

context_parallel_size 在 _validate_linear_cp_mode 执行时已经是 _validate_dynamic_context_parallel 算出的最终值。

6. 三条 provider 路径直接透传

参数名称与 MCore dataclass 字段完全相同，因此不需要 resolver：

raw/non-Bridge：core_transformer_config_from_args(args) 会自动复制 args.linear_cp_mode；
Bridge：把 "linear_cp_mode" 加入现有 bridge_keys override 列表，在 provider.finalize() 前直接写入；
标准 custom provider：如果内部调用 core_transformer_config_from_args()，同样自动获得该字段；完全手写 TransformerConfig(...) 的 custom provider 必须主动传 linear_cp_mode=get_args().linear_cp_mode。

模型构造完成且确认静态/dynamic CP 已启用后，setup_model_and_optimizer() 在 _patch_gdn_for_dynamic_cp() 安装完成的位置读取三条路径最终生成的 MCore config，仅对 GDN 模型由 rank 0 打印一次实际启用配置：

[GDN CP] role=actor linear_cp_mode=chunkwise TP=2 max_CP=4 key_heads=16 value_heads=32

日志读取的是构造完成后的实际 config，而不是仅回显 CLI requested value；默认 headwise（Bridge override 前后值相同、不会触发通用 Override 日志）也会明确打印。CP=1 且未启用 dynamic CP 时不会打印，避免把“配置字段存在”误报成“CP 实际启用”。

7. 删除构造期 monkey patch

文件：relax/backends/megatron/model.py。删除了：

setup_model_and_optimizer 里的调用点（原第 313-314 行附近，if ...: _relax_gdn_cp_config_assert()）；
_relax_gdn_cp_config_assert 函数定义本身（原第 413-468 行，约 56 行）。

替代机制：Stage 1 已让 linear_cp_mode="all_gather" 使用 TP-only 规则构造不可整除几何；Stage 2 的 §6 直接把具体 mode 透传给 MCore，不再需要临时改写 head 数骗过校验。

8. GDN wrapper 重构

8.1 patched `GatedDeltaNet.forward` 已经自己解析动态 CP group 和 mode

Stage 1 backport 的 forward（megatron/core/ssm/gated_delta_net.py）在函数最开头就做了：

base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp
cp_group = resolve_cp_group(base_cp_group, packed_seq_params)
cp_size = cp_group.size() if cp_group is not None else 1

if cp_size == 1:
    ...
elif self.config.linear_cp_mode == "headwise":
    ...
elif self.config.linear_cp_mode == "chunkwise":
    ...
elif self.config.linear_cp_mode == "all_gather":
    raise RuntimeError("linear_cp_mode='all_gather' is implemented by the Relax "
                        "GatedDeltaNet wrapper, not by Megatron. ...")

也就是说：headwise/chunkwise 不需要 Relax wrapper 实现算法或修改状态——wrapper 只解析当前 micro-batch 的 CP size 并做分支判断，随后把这两种模式直接转发给原始 MCore forward；原始 forward 自己从 packed_seq_params.cp_group 解析动态 group/size，并读取构造期固定的 self.config.linear_cp_mode。当配置为 all_gather 且当前 micro-batch 的 cp_size > 1 时，wrapper 改为调用 Relax 实现的 All-gather fallback；CP=1 时不需要任何跨 rank gather，因此仍直接走普通 MCore forward。

8.2 新的 dispatcher

def _patch_gdn_for_dynamic_cp() -> None:
    ...
    _orig_forward = GatedDeltaNet.forward

    def _dcp_gdn_forward(self, hidden_states, attention_mask,
                         inference_context=None, packed_seq_params=None, *args, **kwargs):
        cp_size, cp_group, cp_rank = _resolve_gdn_cp(self, packed_seq_params)
        if cp_size == 1 or self.config.linear_cp_mode != "all_gather":
            return _orig_forward(self, hidden_states, attention_mask,
                                  inference_context, packed_seq_params, *args, **kwargs)
        # 原有 All-gather fallback 主体继续原地执行。
        ...

    GatedDeltaNet.forward = _dcp_gdn_forward
    GatedDeltaNet._dcp_patched = True

和原版Relax 的 _dcp_gdn_forward/_call_orig_with_dynamic_cp/native_ok 相比：

原版Relax（删除）	当前PR（新）
每次 forward 现算 `native_ok = num_key_heads % (tp_size*cp_size)==0`，动态决定走 headwise 还是 all-gather	只读 `self.config.linear_cp_mode`（构造时定死），不再看 runtime head 整除情况
调用原始 forward 前临时 `self.cp_size = cp_size; self.pg_collection.cp = cp_group`，`finally` 还原	直接 `_orig_forward(...)`，不碰 `self.cp_size`/`self.pg_collection.cp`——因为原始 forward 自己会从 `packed_seq_params` 解析动态组（见 8.1）
大段 All-gather fallback 嵌套在 `_dcp_gdn_forward`	保持原位置，只改顶部路由和一处 shared-state mutation，避免无意义的代码移动

9. all-gather fallback：显式 `cp_size_headwise=1`

Stage 1 PR #251 已把 MCore GatedDeltaNet._prepare_qkv_for_gated_delta_rule() 的 cp_size_headwise 参数改为 optional；调用方不传时仍回退读取 self.cp_size，用于兼容 Stage 0。旧 All-gather fallback 因为已经恢复完整序列、没有按 CP 切 heads，只能临时把 self.cp_size 改成 1 再调用该函数。

Stage 2 改为直接使用这个接口：

self._prepare_qkv_for_gated_delta_rule(..., cp_size_headwise=1)

因此可以删除对 self.cp_size 的临时写入，计算语义不变。另外在 All-gather 分支入口增加 assert is_thd：Relax fallback 依赖 THD 的 cu_seqlens 和 zig-zag shard 重组，非 THD 输入直接给出明确错误。

10. 完成标准

项	状态
1. 用户可以通过 CLI 选择三个具体 mode	完成（§4，默认 headwise）
2. Relax 不实现 `auto`/resolver	完成（参数直接使用 MCore 字段名和值）
3. MCore 构造期和 forward 使用同一个静态 mode	完成（无运行时切换）
4. 构造期不再 monkey-patch head 数	完成（§7，函数已删除）
5. headwise/chunkwise 由 MCore 执行	完成（§8.1，dispatcher 直接转发）
6. all-gather 由 Relax fallback 执行	完成（§9）
7. GDN wrapper 内 dynamic CP 只改变 final `PackedSeqParams` group/size	完成（wrapper 不再改 shared state；VLM 外层既有 mutation 不在本 PR 范围内）
8. forward 不创建 group、不改 shared module/config 状态	完成，并由 CPU 状态测试与 §12 真实训练覆盖
9. GDN/layout/loss/recompute 使用同一个最终 metadata	现有 flow 保持不变，本 PR 未改 `data.py`/`loss.py`

11. 测试

11.1 CPU

Stage 1 回归文件 test_gdn_chunkwise_cp_layout.py 共 51 项，其中两项直接覆盖 Chunkwise 长度 fail-fast：

test_rank_indices_reject_lengths_not_divisible_by_two_cp：每条 packed sequence 长度不能被 2×CP 整除时，在布局 route 构造阶段报错。
test_gdn_rejects_packed_lengths_not_divisible_by_cp：cu_seqlens 中任一序列长度不能被 CP 整除时，在 GatedDeltaNet._resolve_cu_seqlens() 报错。

Stage 2 新增 test_gdn_cp_mode_stage2.py 共 19 项

CLI/参数校验：默认 headwise，三个具体值可解析；auto/非法值、chunkwise + --allgather-cp 和 chunkwise + deterministic CP>1 明确报错。
dispatcher 分支（Steps 4-6）：
- CP=1 时三种 mode 都直接转发；
- headwise/chunkwise 在 CP>1 时也直接转发；
- all_gather 且 CP>1 进入原有 Relax fallback；
- 动态组优先于静态组（静态 cp_size=8 的模块跑一个动态 cp_size=1 的 micro-batch，必须走 CP=1 分支而不是被静态组的 8 带偏）；
- headwise/chunkwise 和 CP1 All-gather 路径不修改共享 module/config 状态；
- All-gather 的 THD/inference/deterministic 前置断言在真实张量运算前触发。

本地运行:

Stage 1 layout/config: 51 passed
Stage 2 routing:       19 passed

12. 端到端正确性与性能实验

12.1 统一配置与正确性口径

统一使用 8×H200、Qwen3.5-9B、OpenMathReasoning-mini SFT、TP2、SP on、full recompute 和 --megatron-to-hf-mode bridge；训练数据为 text-only，不包含视觉输入。四方配置为老镜像(旧代码)默认路径，以及候选镜像(本次PR以及升级后的FLA+MCore)强制 headwise、all_gather、chunkwise。每次运行 220 个训练 step，前 20 step warmup，后 200 step 计量；

基础 recipe 使用仓库已有脚本(该脚本默认full recompute，可能是因为all gather必须和recompute一起开；所以看不出cp的显存收益，因此本次测试只注重于实际端到端正确性和速度性能)：

scripts/training/sft/run-qwen3.5-9B-math-dynamic-cp-8xgpu.sh

E1a：不设置 CP 环境变量，使用脚本默认的 --dynamic-context-parallel（max CP=4）。
E1b：设置 CP=4，脚本切换为静态 --context-parallel-size 4。
老镜像使用脚本默认路径；候选实验的临时包装脚本仅在原命令后追加 --linear-cp-mode {headwise,all_gather,chunkwise}，其余 recipe 参数保持一致。

正确性口径：同一 seed 下各模式的有效 token 数必须精确相同，并逐 step 对比训练 loss。E1b seed 1234 的 220 step（step 0–219）中，任意两个配置间的单步 loss 绝对差值最大为 0.000486（headwise vs all_gather），均值为 0.00006–0.00008；差异处于不同 CP 通信/规约顺序引入的 BF16 数值噪声范围。

12.2 E1a：Dynamic CP 训练

CP 在 {1,2,4} 中动态选择。表中吞吐为 3 个 seed 的几何平均，loss 为 3 个 seed 的平均值；每个 seed 均统计 warmup 后的 200 step。

模式	effective throughput	mean train/loss	相对 headwise
老镜像默认	29,111.2 tok/s	≈0.3384	0.978
headwise	29,765.0 tok/s	0.3384	1.000
all_gather	32,829.1 tok/s	0.3384	1.103
chunkwise	26,319.2 tok/s	0.3384	0.884

Correctness：四组均完成 220 step，覆盖 CP1/2/4，无整除、packed length 或 collective 错误；同 seed 有效 token 数精确一致，mean train/loss 差值 <0.002。

结论：all_gather 最快；chunkwise 比 headwise 慢 11.6%。

12.3 E1b：静态 CP4 训练

单 seed 1234，所有 microbatch 均使用 CP4；统计 warmup 后的 200 step。

模式	throughput	mean step time	mean train/loss	相对 headwise
老镜像默认	23,602.6 tok/s	12.25 s	0.3384	1.023
headwise	23,081.4 tok/s	12.53 s	0.3384	1.000
all_gather	24,514.7 tok/s	11.79 s	0.3384	1.062
chunkwise	19,802.6 tok/s	14.60 s	0.3384	0.858

Correctness：四组有效 token 数均为 57,820,637，mean train/loss 均为 0.3384，220 step 全部完成。

结论：完整排序为 all_gather > 老镜像默认 > headwise > chunkwise。

12.4 Profiler 根因

模式	GPU 忙碌率	kernel 总耗时
all_gather	90.1%	9.70 s
headwise	73.6%	12.55 s
chunkwise	56.9%	10.64 s

模式	主要通信	耗时	调用次数
headwise	`all_to_allv`	7,893 ms	32,254
all_gather	`all_gather` + `reduce_scatter`	1,156 ms	720
chunkwise	`all_to_allv` + 状态 `all_gather`	3,777 ms	2,880

chunkwise 的同步热点：

get_thd_context_parallel_rank_indices 中的 cu[0].item()：1,113 ms / 30,720 次。
_resolve_cu_seqlens 中的 [-1].item()：921 ms / 1,920 次。
全部 .item()/sync：3.1 s / 34,148 次，GPU 约 43% 时间空闲。

all_gather 虽重复 CP 倍计算，但使用少量大块 collective；headwise 被大量小 all_to_allv 拖慢；chunkwise 被布局/状态通信，以及每层重复 route 引入的 CPU-GPU 同步拖慢。#5664 可优化该瓶颈，但不在我们这次pr范围内。

12.5 结论

三种候选模式的 GDN forward/backward 与 220-step 训练正确性全部通过：相同有效 token、mean train/loss 差值 <0.002、dynamic CP1/2/4 无报错；
在当前 8×H200、Qwen3.5-9B、TP2/CP4 训练配置下，chunkwise 性能比 headwise 差，不建议设为默认或推荐；默认 headwise 合理。
all_gather 在当前实验下最快，是因为通信实现更高效，而不是省计算或显存；它重复 CP 倍 scan，且 GDN 内部几乎不享受 CP 激活显存节省，长上下文或大 CP 下可能 OOM 或变慢。
chunkwise 的瓶颈是 #3282 基线每层重复 route 和大量 .item() 同步；#5664 的后续优化可能改变结果。

Compatibility layer only: FLA 0.4.1 -> 0.4.2 and a selective backport of Megatron-LM `5139086e` (NVIDIA/Megatron-LM#3282) onto the pinned MCore `85bced0a`. No Relax routing changes -- `relax/` is untouched, `auto` never selects chunkwise, and every existing recipe runs the same path as before. RFC: redai-studio#213, reworked per the 2026-08-05 review decisions: * the GDN CP mode is **static** for the whole process. There is no per-call override; `linear_cp_mode` is read by both the construction-time head check and by `GatedDeltaNet.forward`, and nothing in a forward writes to `self` or the shared config. Dynamic CP varies only `cp_group` / `local_cp_size`. * **v1 depends on #3282 only.** Nothing from the still-open #5664 is included. # ⭐ Feature ## MCore backport (docker/patch/megatron/20260805-85bced0ae.patch) - New `megatron/core/context_parallel_layout.py`, **byte-identical to `5139086e` below the module docstring**: zigzag <-> contiguous THD/SBHD partitions and a single-all-to-all swap between them. The THD swap rebuilds its routing from `cu_seqlens` per call, which is upstream behaviour. - `packed_seq_params.py`: `resolve_cp_group()` only. The dataclass field list is untouched. - `transformer_config.py`: `linear_cp_mode` with headwise `% (tp*cp)` vs chunkwise `% tp` head divisibility. Default is `headwise`, NOT upstream's `chunkwise`, so upgrading the image cannot silently reroute a recipe. `all_gather` is accepted as a third declared value using the TP-only rule, so the declared config equals the resolved `--gdn-cp-mode` rather than declaring one mode while running another. Unknown values, including an unresolved `auto`, assert at construction. - `gated_delta_net.py`: `_resolve_cp_routing()` gives the whole CP group to exactly one of headwise / chunkwise and `None` to the other; validates `local_cp_size == cp_group.size()`; never creates a process group; short circuits on `cp_size == 1` before the mode is read so a CP=1 micro-batch is legal under any declared mode; raises if `all_gather` reaches MCore's forward with cp>1 (the Relax wrapper was not installed). Plus zigzag<->contiguous conversion around conv + scan and `cp_context` for both FLA kernels. - Backwards compatible by construction: `cp_context=` is only passed when chunkwise is active, so with it off the FLA call is byte-identical to before and still works against FLA 0.4.1; `_prepare_qkv_for_gated_delta_rule` gains an optional argument so Relax's all-gather fallback keeps calling it unchanged; both existing Relax GDN fixes are preserved verbatim. ## Dependency - `docker/Dockerfile`: `flash-linear-attention==0.4.2` (first release carrying `fla.ops.cp`), plus build-time capability assertions after the FLA install and after the patch apply. One of them asserts `linear_cp_mode` is **absent** from `GatedDeltaNet.forward`, so re-introducing a per-call override fails the build. --- # ✅ Tests ## tests/backends/megatron/test_gdn_chunkwise_cp_layout.py (45 CPU tests) - Both partitions cover every token exactly once for CP in {1,2,4,8} and are permutations of each other. - MCore's zigzag partition is token-for-token identical to Relax's `slice_with_cp` and `gdn_cp_slice`. - Construction gate: default is `headwise`; chunkwise and `all_gather` use the TP-only head rule; `auto` and other unresolved values are rejected. ## tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py (7 NCCL tests) - FLA `causal_conv1d` / `chunk_gated_delta_rule` under `cp_context` vs no CP, in fp32 and bf16, including `dweight`/`dbias`. - Full `GatedDeltaNet` CP=2 vs CP=1 in fp32 and bf16, with headwise run side by side as the control: in fp32 the two CP algorithms' deviation from CP=1 agrees to 1.00x-1.03x per tensor. - zigzag -> contiguous -> zigzag over a real CP group is token-exact, for packed THD with unequal-length samples and for SBHD. - Illegal combinations fail fast: `all_gather` reaching MCore's forward, mismatched `local_cp_size`, chunkwise + deterministic, chunkwise + inference. - `state_dict` / `sharded_state_dict` keys and shard dims identical across CP=1 / headwise / chunkwise. ## tests/backends/megatron/gdn_cp_numeric_probe.py - Cross-image probe for RFC 3.4-2. Old image vs new image on CP=1 / headwise / all-gather: 41 of 45 tensors bitwise identical and **0 outside tolerance**. The four that differ are at relative RMS 1e-9..6e-7 with cosine 1.0000000000 -- FLA 0.4.2 reorders a few backward reductions. All-gather is 18 of 18 bitwise identical. --- # 📝 Documentation - `docker/patch/megatron/TASK32-BACKPORT.md`: file-level and hunk-level record of what came from `5139086e`, which Relax adaptations were made, which existing Relax GDN fixes are preserved, and what was excluded. Includes two mechanical commands a reviewer can run to confirm the new module matches upstream byte-for-byte and that no #5664 content is present. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

xiaoliang0601 · 2026-08-10T11:59:50Z

有点反直觉。chunkwise 居然比 all_gather 要慢吗？我有时间仔细 review 下。

jambow0320 · 2026-08-10T15:27:32Z

有点反直觉。chunkwise 居然比 all_gather 要慢吗？我有时间仔细 review 下。

嗯嗯，我这边抓了trace，chunkwise大部分耗时都在zigzag、contiguous的布局转换上了，这部分不做优化的话，理论上就是会比all gather慢，有很多cpu端操作和gpu端的小算子

xiaoliang0601 · 2026-08-13T03:36:31Z

你尝试下把 #5664 的核心优化接入下；另外调整下 recompute 和 sequence length，看看有没有正向收益呢？

jambow0320 · 2026-08-13T03:45:40Z

你尝试下把 #5664 的核心优化接入下；另外调整下 recompute 和 sequence length，看看有没有正向收益呢？

okk我试试；不过活动好像明天就截止了，可能一时半会搞不完；我先试试看

jambow0320 · 2026-08-13T13:32:47Z

@xiaoliang0601 Hi~ 我接入了5664的缓存优化，现在chunkwise的性能已经是实验结果最优的了；#273

jambow0320 and others added 4 commits August 5, 2026 20:02

version v2

31b38db

fix headwise resolve cu_seqlens check

88c6a90

relax cp chunkwise intergration

22b1206

jambow0320 requested review from NINGBENZHE and Yangruipis as code owners August 10, 2026 10:23

jambow0320 mentioned this pull request Aug 10, 2026

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

Open

jambow0320 mentioned this pull request Aug 13, 2026

[Task32] feat: fla mcore chunkwise cp stage 3 #273

Open

Navigation Menu

Uh oh!

[Task32] feat: fla mcore chunkwise cp stage 2 - #254

[Task32] feat: fla mcore chunkwise cp stage 2#254
jambow0320 wants to merge 4 commits into
redai-studio:mainfrom
jambow0320:task32-fla-mcore-chunkwise-cp-stage2

jambow0320 commented Aug 10, 2026 •

edited

Loading

Uh oh!

xiaoliang0601 commented Aug 10, 2026

Uh oh!

jambow0320 commented Aug 10, 2026

Uh oh!

xiaoliang0601 commented Aug 13, 2026

Uh oh!

jambow0320 commented Aug 13, 2026

Uh oh!

jambow0320 commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

jambow0320 commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Task 32 第二阶段：Relax 侧 GDN CP 路由接入

1. Summary

2. Background

3. Scope

4. --linear-cp-mode CLI 参数

5. 参数校验

6. 三条 provider 路径直接透传

7. 删除构造期 monkey patch

8. GDN wrapper 重构

8.1 patched GatedDeltaNet.forward 已经自己解析动态 CP group 和 mode

8.2 新的 dispatcher

9. all-gather fallback：显式 cp_size_headwise=1

10. 完成标准

11. 测试

11.1 CPU

12. 端到端正确性与性能实验

12.1 统一配置与正确性口径

12.2 E1a：Dynamic CP 训练

12.3 E1b：静态 CP4 训练

12.4 Profiler 根因

12.5 结论

Uh oh!

xiaoliang0601 commented Aug 10, 2026

Uh oh!

jambow0320 commented Aug 10, 2026

Uh oh!

xiaoliang0601 commented Aug 13, 2026

Uh oh!

jambow0320 commented Aug 13, 2026

Uh oh!

jambow0320 commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jambow0320 commented Aug 10, 2026 •

edited

Loading

4. `--linear-cp-mode` CLI 参数

8.1 patched `GatedDeltaNet.forward` 已经自己解析动态 CP group 和 mode

9. all-gather fallback：显式 `cp_size_headwise=1`