执行摘要

• 一句话：重构 --save-debug-train-data 为单文件转储并支持 CP 恢复
• 推荐动作：值得精读。重点看 restore_context_parallel_fields_to_cpu 的逐 sample gather 与显存控制、policy_loss_function 中的捕获钩子设计，以及 _build_dump_payload 的排序与布局。对需要扩展训练数据转储或理解 Megatron 后端 CP/PP 数据流转的工程师有直接参考价值。

功能与动机

PR body 为空，动机可从代码注释和提交信息还原。旧实现每 rank 各写一个文件，样本顺序与 rollout dump 不对齐，无法逐样本比对；CP 开启时 response 级字段只有本 rank 局部片段，落盘数据不可直接用；当配置跳过单独 log-prob 重算（can_reuse_log_probs_in_loss / use_rollout_logprobs）时转储拿不到 per-sample log_probs，需要额外 forward。新实现希望在不多跑 forward、不显著抬高显存峰值的前提下，产出与 rollout dump 严格对齐的规范化 train dump。

实现拆解

1. 转储核心迁移与重建：新增 slime/backends/megatron_utils/train_dump_utils.py，删除 slime/utils/train_dump_utils.py。新的 save_debug_train_data 只由“最后一个 PP stage + TP rank 0 + CP rank 0”的 writer rank 写盘；_build_dump_payload 把 DP gather 到的各 rank 载荷整理成与 rollout dump 对齐的结构：samples 按 partition 优先排序、sample_indices 兜底；micro_batch_indices、num_microbatches、global_batch_sizes 等布局字段挂在 dp_shards 键下，避免复制 per-sample 张量。

2. CP 字段逐步恢复与显存控制：restore_context_parallel_fields_to_cpu 遍历 _CONTEXT_PARALLEL_FIELDS（共 10 个 response 级字段，如 rollout_log_probs、values、advantages、returns）逐 sample 调用 gather_tensor(value, total_length, response_length) 恢复完整序列；keep_restored 只对 writer rank 生效，其余 rank 跑完集合通信即丢弃结果，每次 gather 后立刻 del full_value，把额外设备显存峰值压到“一个 response 张量”以内。

3. 训练期 log-prob 捕获：slime/backends/megatron_utils/loss.py 新增模块级 _LOG_PROB_CAPTURE 及 enable_log_prob_capture、drain_captured_log_probs、_maybe_capture_log_probs；policy_loss_function 在 torch.cat 重绑定 log_probs 之前调用捕获钩子，按全局 partition 位置保存 CP-local 张量，“首次出现优先”保证多 step 训练保留 old-policy 值。actor.py 的 train_actor 在开启 dump 且 rollout_data 缺少 log_probs 时启用捕获，train 结束后以本地 partition 顺序还原注入 rollout_data["log_probs"]。

4. 数据契约与配置联动：slime/utils/data.py 在 dump 模式下向训练数据注入 partition 键；model.py 配合格式版本控制；arguments.py 微调参数语义；dump 文件改为 format_version=2，单文件同时含 samples 和 dp_shards 两个顶层键。

5. 测试、CI 与文档：新增 tests/test_train_dump.py（真实集合通信的 CP 恢复、DP 分片单文件写入、无 CP 归一化格式）与 8 GPU e2e tests/test_qwen2.5_0.5B_debug_train_dump_e2e.py（TP=2、PP=2、CP=2、DP=1，按 rollout_position 连接两份 dump 并逐样本比对 rollout_log_probs）；pr-test.yml 及其模板注册新测试，docs/en/developer_guide/debug.md 补充新格式说明。

关键文件：

• slime/backends/megatron_utils/train_dump_utils.py（模块 转储工具；类别 source；类型 dependency-wiring；符号 _to_cpu, restore_context_parallel_fields_to_cpu, _is_per_sample, _build_dump_payload）: 转储核心实现，从 slime/utils 迁入并重写为单文件格式；包含 CP 字段恢复、payload 构建与 writer 选择逻辑，是本 PR 的主体。
• slime/backends/megatron_utils/loss.py（模块 损失函数；类别 source；类型 core-logic；符号 enable_log_prob_capture, drain_captured_log_probs, _maybe_capture_log_probs）: 在策略损失主路径新增可选 log-prob 捕获钩子，使 dump 在复用训练 forward 时零额外开销拿到 per-sample log-probs。
• slime/backends/megatron_utils/actor.py（模块 训练器；类别 source；类型 dependency-wiring；符号 train_actor）: 接线点：根据 save_debug_train_data 与 rollout_data 是否含 log_probs 决定是否开启捕获，train 后按 partition 还原并注入。
• tests/test_train_dump.py（模块 单元测试；类别 test；类型 test-coverage；符号 _patch_single_dp_writer, _restore_context_parallel_worker, _single_file_dump_worker, test_restore_context_parallel_fields_with_real_collective）: 覆盖新转储核心：真实集合通信 CP 恢复、DP 分片单文件写入、无 CP 归一化格式。
• tests/test_qwen2.5_0.5B_debug_train_dump_e2e.py（模块 端到端测试；类别 test；类型 test-coverage；符号 prepare, _train_args, _verify, execute）: 8 GPU 端到端验证：TP/PP/CP 全开，连接 rollout 与 train dump 逐样本比对 rollout_log_probs。
• slime/utils/train_dump_utils.py（模块 转储工具；类别 source；类型 deletion；符号 save_debug_train_data）: 旧实现被删除，逻辑迁移至 megatron_utils 下的新文件。

关键符号：save_debug_train_data, restore_context_parallel_fields_to_cpu, _build_dump_payload, _to_cpu, _is_per_sample, enable_log_prob_capture, drain_captured_log_probs, _maybe_capture_log_probs

关键源码片段

slime/backends/megatron_utils/loss.py

在策略损失主路径新增可选 log-prob 捕获钩子，使 dump 在复用训练 forward 时零额外开销拿到 per-sample log-probs。

# 训练 forward 中产生的 per-sample log-probs 快照，按全局 rollout position 索引。
# 只在需要 dump 且复用训练 forward 的 log-probs 时启用，避免额外 forward。
_LOG_PROB_CAPTURE: "dict[int, torch.Tensor] | None" = None
​
​
def enable_log_prob_capture() -> None:
    """开启捕获，必须在 train 之前调用。"""
    global _LOG_PROB_CAPTURE
    _LOG_PROB_CAPTURE = {}
​
​
def drain_captured_log_probs() -> "dict[int, torch.Tensor]":
    """取走捕获结果并停止捕获；未捕获时返回空 dict。"""
    global _LOG_PROB_CAPTURE
    captured = _LOG_PROB_CAPTURE or {}
    _LOG_PROB_CAPTURE = None
    return captured
​
​
def _maybe_capture_log_probs(batch: RolloutBatch, log_probs: list[torch.Tensor]) -> None:
    """在 policy_loss_function 中快照 per-sample CP-local log-probs。
​
    仅在捕获开启且 batch 携带 partition 时工作；partition 只在 dump 模式下
    被注入训练数据。每个 position 保留首次出现的值（old-policy），多 step
    训练因此不会覆盖初始策略的 log-probs。
    """
    if _LOG_PROB_CAPTURE is None:
        return
    positions = batch.get("partition")
    if not positions:
        return
    for position, log_prob in zip(positions, log_probs, strict=True):
        if position not in _LOG_PROB_CAPTURE:
            _LOG_PROB_CAPTURE[position] = log_prob.detach().clone()

评论区精华

本 PR 没有 review 评论或讨论线程（comments_count = 0、review_comments_count = 0）。设计取舍体现在代码注释与提交信息中：单文件格式中 dp_shards 布局单独存放，是为了与 rollout debug dump 的扁平 samples 视图对齐而不复制 per-sample 张量；writer 只由“last PP stage + TP0 + CP0”承担，e2e 测试专门覆盖该选取逻辑；keep_restored 仅在 writer 上保留 CPU 值并配合 del full_value 控制显存峰值，是最值得借鉴的写法。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 核心训练路径改动：loss.py 的 policy_loss_function 是训练主路径，_maybe_capture_log_probs 的新分支必须保证默认关闭时零开销、零行为变化；注释明确要求钩子在 torch.cat 重绑定 log_probs 之前执行，后续维护需保持此顺序。
  2. 多 rank 状态管理：模块级 _LOG_PROB_CAPTURE 依赖 actor.py 的对称启停；若 train 异常退出未调用 drain_captured_log_probs，可能残留旧状态，建议补充异常路径清理。
  3. 新增 CP 集合通信：restore_context_parallel_fields_to_cpu 对每个 CP 字段逐个 sample gather，字段多时通信次数线性增加；虽有 del 控制显存，通信延迟对训练循环仍可能有扰动，目前仅在 dump 模式启用。
  4. 数据格式契约变更：dump 产物从“每 rank 一个文件”变为“单文件 + format_version=2”，依赖旧格式的调试脚本需要迁移；partition 键只在 dump 模式注入，正常训练数据契约不变。
  5. e2e 覆盖有限：8 GPU e2e 覆盖 TP=2/PP=2/CP=2/DP=1，未覆盖 DP>1 与 CP 同时存在的组合（单元测试用 4 进程模拟 DP=2/CP=2）。
     - 影响：影响范围集中在 Megatron 后端训练调试链路：使用 --save-debug-train-data 的用户会拿到格式变更后的产物（单文件、format_version=2），需要相应调整调试脚本；新 log-prob 捕获仅在开启 dump 且复用训练 forward 时生效，对正常运行无感知；CI 增加 8 GPU e2e 测试成本；团队后续可以用同一份规范化 dump 直接对齐 rollout 与 train 数据，降低排查策略/奖励不一致的难度。
     - 风险标记：核心训练路径改动, 数据格式契约变更, 新增 CP 集合通信, 多 rank 状态管理

关联脉络

• PR #2247 fix: forward dual-clip PPO epsilon: 同改 slime/backends/megatron_utils/loss.py，都在训练损失主路径上做配置/数据接线。
• PR #2234 fix: pair --log-correct-samples rewards with the DP-local samples: 同改 slime/utils/data.py，都处理 DP 切分下样本与张量的对齐问题。
• PR #2205 perf: vectorize REINFORCE++ discounted returns: 同属训练工具链数值逻辑重构，并配套新增测试，反映该区域的持续演进。