Refactor --save-debug-train-data

.github/workflows/pr-test.yml

@@ -228,6 +228,10 @@ jobs:
               "num_gpus": 8,
               "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"
             },
+            {
+              "num_gpus": 8,
+              "test_file": "test_qwen2.5_0.5B_debug_train_dump_e2e.py"
+            },
             {
               "num_gpus": 8,
               "test_file": "test_qwen2.5_0.5B_opd_sglang.py"
@@ -589,6 +593,10 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_cp_utils.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_train_dump.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_metric_report.py"
@@ -988,6 +996,10 @@ jobs:
               "num_gpus": 8,
               "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"
             },
+            {
+              "num_gpus": 8,
+              "test_file": "test_qwen2.5_0.5B_debug_train_dump_e2e.py"
+            },
             {
               "num_gpus": 8,
               "test_file": "test_qwen2.5_0.5B_opd_sglang.py"

.github/workflows/pr-test.yml.j2

@@ -17,6 +17,7 @@
     {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8},
     {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8},
     {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8},
+    {'test_file': 'test_qwen2.5_0.5B_debug_train_dump_e2e.py', 'num_gpus': 8},
     {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8},
     {'test_file': 'test_qwen3_4B_external_pd.py', 'num_gpus': 6},
     {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4},
@@ -73,6 +74,7 @@
         {'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0},
         {'test_file': 'test_dp_schedule.py', 'num_gpus': 0},
         {'test_file': 'test_cp_utils.py', 'num_gpus': 0},
+        {'test_file': 'test_train_dump.py', 'num_gpus': 0},
         {'test_file': 'test_metric_report.py', 'num_gpus': 0},
         {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0},
         {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0},

docs/en/developer_guide/debug.md

@@ -50,6 +50,12 @@ Specifically, slime currently provides the following parameters for separate deb
 
     When enabled, data will be loaded from `args.load_debug_rollout_data.format(rollout_id=rollout_id)`, and SGLang will not be initialized (automatically setting `debug_train_only=True`). This method allows you to fix the input for the training part to tune it, for example, by switching between different parallelization strategies.
 
+5.  `--save-debug-train-data /your/saved/debug/train_{rollout_id}.pt`
+
+    Saves one train-side file per rollout. Only the last Pipeline Parallel stage and Tensor Parallel rank 0 participate. They restore response-token fields such as `log_probs`, `ref_log_probs`, `values`, `advantages`, `returns`, `kl`, and `entropy` across Context Parallel ranks. Context Parallel rank 0 moves each restored tensor to CPU immediately, so complete tensors do not accumulate on the GPU, and then gathers the distinct Data Parallel shards to one writer.
+
+    The version-2 payload mirrors the rollout debug dump: a top-level `samples` list holds one dict per training sample (`sample_index`, `data_parallel_rank`, and its per-sample fields such as `tokens`, `log_probs`, `advantages`), sorted by `sample_index` so it lines up one-to-one with the rollout dump's `samples` (join on `sample_index` ↔ the rollout side's `index`). A parallel `dp_shards` key preserves the DP/micro-batch layout — each entry records `rank`, `data_parallel_rank`, that shard's `sample_indices`, and the DP-local schedule (`micro_batch_indices`, `num_microbatches`, `global_batch_sizes`) — without duplicating any per-sample tensor. Whole-batch fields such as `raw_reward` are stored once at the top level. If any sample lacks a `sample_index` (custom rollouts that build fresh `Sample` objects leave it `None`), the samples stay in DP-gather order and a warning is logged. With or without CP, response-token fields use the same full-response format. In configs that skip the separate actor log-prob recompute (`can_reuse_log_probs_in_loss` or `--use-rollout-logprobs`), the actor `log_probs` are snapshotted from the training forward itself (keyed by rollout position, at no extra forward), so the dump still carries them.
+
 ## INT4 / Compressed-Tensors Quantization Checkpoint Issues
 
 When using INT4-quantized models (e.g., `compressed-tensors` with `W4A16`), the checkpoint's `config.json` contains a `quantization_config.ignore` list that specifies which parameters should **not** be quantized. During online weight updates (Megatron → SGLang), slime also reads this ignore list to decide which parameters to INT4-quantize. An incorrect ignore list can cause silent errors:

docs/zh/developer_guide/debug.md

@@ -48,6 +48,12 @@ slime 支持将训练部分和推理部分分开进行调试，从而实现：
 
    开启后，会从 `args.load_debug_rollout_data.format(rollout_id=rollout_id)` 来加载数据，并且不会初始化 sglang（自动设置 `debug_train_only=True`）。可以以这种方式来固定训练部分的输入，对训练部分进行调优，例如切换各种并行。
 
+4. `--save-debug-train-data /your/saved/debug/train_{rollout_id}.pt`
+
+   每个 rollout 只保存一个训练侧文件。只有 Pipeline Parallel 最后一级和 Tensor Parallel rank 0 参与：跨 Context Parallel rank 逐个还原 `log_probs`、`ref_log_probs`、`values`、`advantages`、`returns`、`kl`、`entropy` 等 response-token tensor；Context Parallel rank 0 会将每个完整 tensor 立即搬到 CPU，避免它们在显存中累计，最后再把不同的 Data Parallel shard 汇总给一个 writer。
+
+   version 2 payload 对标 rollout debug dump：顶层 `samples` 列表每项是一个训练样本的 dict（含 `sample_index`、`data_parallel_rank` 以及 `tokens`、`log_probs`、`advantages` 等 per-sample 字段），并按 `sample_index` 排序，从而和 rollout dump 的 `samples` 一一对齐（用 `sample_index` ↔ rollout 侧的 `index` 来 join）。并列的 `dp_shards` key 保留 DP/micro-batch 排布——每项记录 `rank`、`data_parallel_rank`、该分片的 `sample_indices`，以及 DP-local 调度（`micro_batch_indices`、`num_microbatches`、`global_batch_sizes`）——且不重复存储任何 per-sample tensor。`raw_reward` 等整批字段在顶层只存一份。若某些样本没有 `sample_index`（自定义 rollout 新建 `Sample` 时会是 `None`），则 samples 保持 DP-gather 顺序并打印一条 warning。开启或关闭 CP 时，response-token 字段都是相同的完整 response 格式。在跳过 actor log-prob 单独重算的配置下（`can_reuse_log_probs_in_loss` 或 `--use-rollout-logprobs`），actor 的 `log_probs` 会直接从训练前向里快照下来（按 rollout position 归位，无额外前向），所以 dump 里依然会带上它。
+
 ## INT4 / Compressed-Tensors 量化 Checkpoint 问题
 
 使用 INT4 量化模型（如 `compressed-tensors` 的 `W4A16`）时，checkpoint 的 `config.json` 中有一个 `quantization_config.ignore` 列表，指定哪些参数**不**做量化。在线权重更新（Megatron → SGLang）时，slime 也会读取这个 ignore list 来决定哪些参数需要 INT4 量化。ignore list 不正确会导致静默错误：

slime/backends/megatron_utils/actor.py

@@ -13,7 +13,6 @@
 from transformers import AutoConfig, AutoTokenizer
 
 from slime.ray.train_actor import TrainRayActor
-from slime.utils import train_dump_utils
 from slime.utils.data import process_rollout_data
 from slime.utils.distributed_utils import get_gloo_group
 from slime.utils.logging_utils import init_tracking
@@ -31,12 +30,19 @@
 
 from ...utils.profile_utils import TrainProfiler
 from ...utils.tensor_backper import TensorBackuper
+from . import train_dump_utils
 from .checkpoint import load_checkpoint
 from .cp_utils import prepare_routed_experts_for_routing_replay, slice_log_prob_with_cp
 from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data
 from .hf_checkpoint_saver import save_hf_model_to_path
 from .initialize import init, is_megatron_main_rank
-from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values
+from .loss import (
+    compute_advantages_and_returns,
+    drain_captured_log_probs,
+    enable_log_prob_capture,
+    get_log_probs_and_entropy,
+    get_values,
+)
 from .model import forward_only, initialize_model_and_optimizer, save, train
 from .update_weight.common import named_params_and_buffers
 from .update_weight.update_weight_from_disk import UpdateWeightFromDisk
@@ -508,6 +514,13 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
             # Train
             if self.args.use_routing_replay:
                 os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward"
+            # When dumping train debug data but the actor log_probs were not
+            # recomputed separately (can_reuse_log_probs_in_loss / use_rollout_logprobs),
+            # snapshot them from the training forward so the dump still carries
+            # per-sample log_probs — at no extra forward pass.
+            capture_log_probs = self.args.save_debug_train_data is not None and "log_probs" not in rollout_data
+            if capture_log_probs:
+                enable_log_prob_capture()
             with timer("actor_train"):
                 train(
                     rollout_id,
@@ -518,6 +531,14 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
                     num_microbatches,
                     global_batch_sizes,
                 )
+            if capture_log_probs:
+                captured = drain_captured_log_probs()
+                # `captured` is non-empty only on the last PP stage running a loss
+                # that snapshots log_probs (policy_loss), and then covers every
+                # local sample. Key it by this rank's `partition` to land in local
+                # sample order; skip otherwise (nothing to place).
+                if captured:
+                    rollout_data["log_probs"] = [captured[pos] for pos in rollout_data["partition"]]
 
             self.prof.step(rollout_id=rollout_id)
 

slime/backends/megatron_utils/loss.py

@@ -37,6 +37,49 @@
 )
 
 
+# Optional capture of per-sample policy log-probs computed during the training
+# forward. Used only when dumping train debug data in configs that skip the
+# separate log-prob recompute (can_reuse_log_probs_in_loss / use_rollout_logprobs):
+# the values are identical to a separate compute_log_prob pass, so we snapshot
+# them here at no extra forward. Keyed by GLOBAL rollout position so the writer
+# can put them back in sample order regardless of pipeline/microbatch order.
+_LOG_PROB_CAPTURE: "dict[int, torch.Tensor] | None" = None
+
+
+def enable_log_prob_capture() -> None:
+    """Start capturing training-forward log-probs (call before ``train``)."""
+    global _LOG_PROB_CAPTURE
+    _LOG_PROB_CAPTURE = {}
+
+
+def drain_captured_log_probs() -> "dict[int, torch.Tensor]":
+    """Return captured ``{rollout_position: cp-local log_probs}`` and stop capturing."""
+    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:
+    """Snapshot per-sample CP-local ``log_probs`` keyed by global rollout position.
+
+    No-op unless :func:`enable_log_prob_capture` is active and the micro-batch
+    carries ``partition`` (only added to the training keys when dumping). First
+    occurrence per position wins, so multi-step training keeps the initial
+    (old-policy) values. ``log_probs`` here is the per-sample list, in the same
+    order as ``batch['partition']`` (both indexed by this micro-batch's
+    ``micro_batch_indices``).
+    """
+    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()
+
+
 def get_rollout_top_p_logprob_kwargs(args: Namespace, batch: dict[str, Any]) -> dict[str, Any]:
     if args.rollout_top_p == 1.0:
         return {}
@@ -935,6 +978,10 @@ def policy_loss_function(
     )
 
     log_probs = log_probs_and_entropy["log_probs"]
+    # Snapshot the per-sample policy log-probs for the train debug dump (no-op
+    # unless capture is enabled). Must run before the torch.cat below rebinds
+    # `log_probs` to a single concatenated tensor.
+    _maybe_capture_log_probs(batch, log_probs)
     if not args.use_rollout_logprobs and not old_log_probs:
         old_log_probs = [log_prob.detach() for log_prob in log_probs]
     train_log_probs_for_tis = batch.get("log_probs")

slime/backends/megatron_utils/model.py

@@ -593,6 +593,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                     "rollout_log_probs",
                     "teacher_log_probs",
                     "rollout_mask_sums",
+                    # Only present when dumping train debug data; lets the loss
+                    # snapshot each sample's log_probs keyed by rollout position.
+                    *(["partition"] if args.save_debug_train_data is not None else []),
                 ],
             ),
             args.data_pad_size_multiplier,