Support DSV4 IndexCache training, distillation, and recompute

[chenbong]

Summary

• add DSV4 IndexCache F/S pattern configuration, validation, producer/reuse flow, and non-uniform pattern support
• carry explicit IndexCache state across PP/CP boundaries, including FSF handoff
• support reuse-only and multi-layer distillation with full/uniform recompute
• preserve producer gradients through schedule and outer pipeline backward boundaries, including missing or released metadata cases
• add training safeguards for frozen/indexer-only paths, custom backward behavior, debug visibility, and invalid state
• keep checkpoint routing-table dtype compatible with the current training and resume path

Baseline

• base branch: PaddlePaddle/PaddleFleet:release/0.3
• validated base commit: 4dd81f92437e120208e7814267622617add3054b
• current head commit: 10fd153843c6c14a7cae7a72084fc86c16feafa8

The upstream release/0.3 branch has advanced by three commits since this validated base. This draft intentionally preserves the tested history while the change is reviewed.

Scope

• 10 commits
• 13 changed files
• 2,353 insertions and 271 deletions

The stack is intentionally left unsquashed in this draft so the evolution from basic F/S reuse through recompute and PP-gradient fixes remains reviewable.

Validation

• git diff --check 4dd81f9..10fd153
• focused IndexCache PP runtime suite: 6 passed
• downstream single-node 4/8-GPU training regressions, including a 36-step soak
• downstream 16-node/64-GPU two-step RL training and full-state checkpoint resume validation

Upstream synchronization note

GitHub currently reports one content conflict with the three newer release/0.3 commits:

• src/paddlefleet/transformer/transformer_layer.py

It is intentionally left unresolved in this Draft so the validated 10fd153 stack remains unchanged. The conflict should be resolved and the relevant validation repeated before this PR is made ready for final submission.

Review status

This remains a Draft PR for inspecting the complete current integration. Commit cleanup, upstream-base synchronization, and any requested split into smaller PRs are intentionally deferred until the review direction is clear.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

chenbohong seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-13 20:29:27

📋 Review 摘要

PR 概述：为 DSV4 CSA IndexCache 增加 F/S pattern、PP/CP 状态传递、distill/recompute 和相关保护。
变更范围：transformer/csa_attention.py、transformer_layer.py、PP runtime patch、TransformerConfig、MoE/FP8、单测。
影响面 Tag：Transformer Models MoE Fusions Pipeline Parallel Context Parallel

问题

级别	文件	概述
🔴 Bug	src/paddlefleet/transformer/csa_attention.py:1789	MTP ratio=4 层的 layer id 映射与 IndexCache pattern 映射不一致，启用 pattern 后运行时找不到 C4 layer
🟡 建议	tests/single_card_tests/test_indexcache_pp_runtime_patch.py:208	新增 test 文件包含本地 __main__ 入口
🟡 建议	PR-level	本 PR 变更量较大（13 文件 / 2624 行），建议拆分以降低审查和回归风险

历史 Findings 修复情况

Finding	问题	状态
F1	indexcache_multi_layer_distill producer loss state 梯度被断开	✅ 已修复
F2	_indexcache_tensor_to_int() 在 forward 热路径同步 Tensor 到 Python int	⚠️ 仍存在

📝 PR 规范检查

标题缺少规范 Tag，描述未使用 PaddleFleet §D2 模板。

标题建议（可直接复制）：

• [Distributed Strategy] Support DSV4 IndexCache training handoff

PR 描述建议（点击展开，可直接复制）

### PR Category
<!-- One of [ User Experience | Execute Infrastructure | Operator Mechanism | Custom Device | Performance Optimization | Distributed Strategy | Parameter Server | Communication Library  | Environment Adaptation ] -->
Distributed Strategy
### PR Types
<!-- One of [ New features | Bug fixes | Improvements | Performance | BC Breaking | Deprecations | Docs | Devs | Not User Facing | Security | Others ] -->
New features
### Description
<!-- Describe what you’ve done -->
Support DSV4 IndexCache training F/S pattern and state handoff for ratio=4 CSA indexer layers. Add PP/CP IndexCache state propagation, multi-layer distillation plumbing, stop_gradient safeguards for CSA sparse attention/indexer losses, and config validation for `index_topk_pattern` and `indexcache_multi_layer_distill`.

总体评价

⚠️ 本 PR 变更量较大（13 文件 / 2624 行），建议拆分以降低审查难度和合入风险。

建议拆分方案：

• PR 1: IndexCache 状态结构与 PP runtime patch — src/paddlefleet/indexcache_pp_runtime_patch.py, src/paddlefleet/transformer/indexcache_state.py, tests/single_card_tests/test_indexcache_pp_runtime_patch.py
• PR 2: CSA/DSV4 F/S pattern、distill 与 recompute — src/paddlefleet/transformer/csa_attention.py, src/paddlefleet/transformer/dsv4_hybrid_attention.py, src/paddlefleet/transformer/transformer_layer.py, src/paddlefleet/transformer/transformer_config.py
• PR 3: MoE/FP8/hash router 与外围兼容改动 — src/paddlefleet/transformer/moe/*.py, src/paddlefleet/models/gpt/gpt_embedding.py, src/paddlefleet/fusions/csa_sparse_attn.py

本轮按风险优先审查了 IndexCache 状态流、PP dict/tuple 转换、recompute 输出展平、CP attention 分支、TransformerConfig 校验、MoE/FP8 触达 hunk 和新增测试；未全量覆盖所有数值组合。当前环境缺少 paddle，因此未运行单测。

[PaddlePaddle-bot]

🔴 Bug _indexcache_c4_layers() 直接把 csa_compress_ratios 的下标加 num_empty_layers_add_in_head 当成运行时 layer id，但 MTP 层在 CompressedSparseAttention.__init__ 中会被改成 num_hidden_layers + 1 + mtp_idx。因此当 csa_compress_ratios[num_hidden_layers + mtp_idx] == 4 且配置了 index_topk_pattern 时，配置校验会把该 MTP C4 层计入 pattern 长度，而运行到该层时 _indexcache_next_c4_action() 会因为 self.layer_number not in c4_layers 直接抛错，IndexCache reuse/distill 无法用于 MTP C4 层。

建议修复方式：让这里的 C4 layer id 生成规则与 CompressedSparseAttention.__init__ 完全一致；例如对 layer_idx < num_hidden_layers 仍使用 layer_idx + num_empty_layers_add_in_head，对 MTP 部分使用 num_hidden_layers + 1 + mtp_idx。如果暂不支持 MTP，则在 TransformerConfig.__post_init__ 中显式拒绝 MTP 部分含 ratio=4 的 index_topk_pattern 配置。

[PaddlePaddle-bot]

🟡 建议 新增 test 文件包含 if __name__ == "__main__" 本地入口，命中 PaddleFleet checklist §C 的 test 环境依赖信号。CI 应通过测试发现机制执行该 unittest.TestCase，保留脚本入口容易让本地直接运行路径和 CI 路径分叉。

建议修复方式：删除底部入口，保持测试只由 pytest/unittest discovery 运行。