Make the tensor dumper fault-tolerance aware

miles/backends/megatron_utils/actor.py

@@ -261,6 +261,7 @@ def compute_log_prob(
         self,
         data_iterator: list[DataIterator],
         num_microbatches: list[int],
+        rollout_id: int,
         store_prefix: str = "",
     ) -> dict[str, list[torch.Tensor]]:
 
@@ -271,6 +272,7 @@ def compute_log_prob(
                 self.model,
                 data_iterator,
                 num_microbatches,
+                rollout_id=rollout_id,
                 store_prefix=store_prefix,
             )
 
@@ -306,6 +308,7 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
                 self.model,
                 data_iterator,
                 num_microbatches,
+                rollout_id=rollout_id,
             )
         )
 
@@ -359,6 +362,7 @@ def train_actor(
                         self.compute_log_prob(
                             data_iterator,
                             num_microbatches,
+                            rollout_id=rollout_id,
                             store_prefix="ref_",
                         )
                     )
@@ -370,6 +374,7 @@ def train_actor(
                         self.compute_log_prob(
                             data_iterator,
                             num_microbatches,
+                            rollout_id=rollout_id,
                             store_prefix="teacher_",
                         )
                     )
@@ -385,6 +390,7 @@ def train_actor(
                         self.compute_log_prob(
                             data_iterator,
                             num_microbatches,
+                            rollout_id=rollout_id,
                             store_prefix="",
                         )
                     )

miles/backends/megatron_utils/model.py

@@ -213,6 +213,7 @@ def forward_only(
     model: Sequence[DDP],
     data_iterator: Sequence[DataIterator],
     num_microbatches: Sequence[int],
+    rollout_id: int,
     store_prefix: str = "",
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
@@ -226,13 +227,16 @@ def forward_only(
         model: Sequence of DDP-wrapped model chunks.
         data_iterator: Iterable(s) yielding batches for inference.
         num_microbatches: Number of microbatches per rollout step.
+        rollout_id: Rollout identifier (selects the per-rollout dump subdirectory).
         store_prefix: Prefix to prepend to stored output keys.
 
     Returns:
         Aggregated outputs keyed by ``store_prefix + key``.
     """
 
-    dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_ONLY)
+    dumper_phase_util = DumperMegatronUtil(
+        args, model, DumperPhase.FWD_ONLY, rollout_id=rollout_id, store_prefix=store_prefix
+    )
 
     # reset data iterator
     for iterator in data_iterator:
@@ -278,6 +282,7 @@ def forward_step(
         packed_seq_params = get_packed_seq_params(batch, args)
         total_lengths = batch["total_lengths"]
         response_lengths = batch["response_lengths"]
+
         output_tensor = model(
             input_ids=tokens,
             position_ids=None,
@@ -372,7 +377,7 @@ def train_one_step(
         Reduced loss dictionary (last stage only) and gradient norm for logging.
     """
     args = get_args()
-    dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD)
+    dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=rollout_id)
     disable_optimizer = args.debug_disable_optimizer or optimizer is None
 
     # Set grad to zero.

miles/utils/dumper_utils.py

@@ -14,6 +14,9 @@
 import torch.distributed as dist
 from sglang.srt.debug_utils.dumper import DumperConfig, _get_rank, dumper
 
+from miles.backends.training_utils.parallel import get_parallel_state
+from miles.utils.environ import enable_experimental_ft_trainer
+
 logger = logging.getLogger(__name__)
 
 
@@ -78,10 +81,21 @@ async def configure_sglang(args: Namespace) -> None:
 
 
 class DumperMegatronUtil:
-    def __init__(self, args: Namespace, model: Sequence[torch.nn.Module], phase: DumperPhase) -> None:
+    def __init__(
+        self,
+        args: Namespace,
+        model: Sequence[torch.nn.Module],
+        phase: DumperPhase,
+        *,
+        rollout_id: int,
+        store_prefix: str = "",
+    ) -> None:
         self.phase = phase
+        self.rollout_id = rollout_id
         self.overrides = _get_phase_override_configs(args, phase)
-        self.enabled = self._configure(args, phase, self.overrides)
+        self.enabled = self._configure(
+            args, phase=phase, rollout_id=rollout_id, store_prefix=store_prefix, overrides=self.overrides
+        )
         if self.enabled:
             dumper.register_non_intrusive_dumper(self._extract_model(model))
 
@@ -96,10 +110,15 @@ def finalize(self, model: Sequence[torch.nn.Module]) -> None:
             return
 
         extracted_model = self._extract_model(model)
+        get_grad: Callable[[torch.nn.Parameter], torch.Tensor | None] | None = None
         if self.phase is DumperPhase.FWD_BWD and self.overrides.get("enable_model_grad"):
             _log_model_grad_coverage(extracted_model)
+            if enable_experimental_ft_trainer():
+                get_grad = _build_full_grad_getter(extracted_model)
 
-        dumper.dump_model(extracted_model)
+        # Weights/grads are a once-per-rollout end-state, so pin them to step 0 instead of
+        # the running per-microbatch step.
+        dumper.dump_model(extracted_model, get_grad=get_grad, step=0)
         dumper.step()
         dumper.configure(enable=False)
 
@@ -111,25 +130,97 @@ def _extract_model(model: Sequence[torch.nn.Module]) -> torch.nn.Module:
         return model[0]
 
     @staticmethod
-    def _configure(args: Namespace, phase: DumperPhase, overrides: dict[str, Any] | None = None) -> bool:
+    def _configure(
+        args: Namespace,
+        *,
+        phase: DumperPhase,
+        rollout_id: int,
+        store_prefix: str = "",
+        overrides: dict[str, Any] | None = None,
+    ) -> bool:
         if overrides is None:
             overrides = _get_phase_override_configs(args, phase)
         if not overrides.get("enable"):
             return False
 
+        exp_name = f"{phase.value}/{store_prefix}rollout_{rollout_id}"
         merged = {
             "dir": str(_get_dir(args)),
-            "exp_name": phase.value,
+            "exp_name": exp_name,
+            "enable_output_console": False,
             **overrides,
         }
 
+        # Only write dump files on effective DP rank 0 (covers both intra-DP
+        # and indep-DP). Other DP ranks still participate in dumper collectives
+        # (barrier, broadcast, allgather) but don't produce output files.
+        # TODO: optimize — non-DP-rank-0 ranks currently run full dumper logic
+        # (forward hooks, model iteration) without producing output.
+        if get_parallel_state().intra_dp.rank != 0:
+            merged["enable_output_file"] = False
+            merged["enable_output_console"] = False
+
         full_config = DumperConfig(**merged)
         dumper.reset()
+        # Wipe the whole phase dir only at run start (rollout 0). Gating on a
+        # per-process latch instead would make a respawned process re-wipe the
+        # phase dir mid-run, deleting dumps already written by surviving cells.
+        if rollout_id == 0:
+            _cleanup_dump_dir(Path(merged["dir"]) / phase.value)
         _cleanup_dump_dir(Path(merged["dir"]) / merged["exp_name"])
+        _barrier_after_dump_dir_cleanup()
         dumper.configure(**dataclasses.asdict(full_config))
         return True
 
 
+def _build_full_grad_getter(
+    model_chunk: torch.nn.Module,
+) -> Callable[[torch.nn.Parameter], torch.Tensor | None]:
+    """Build get_grad(param): all-gather distributed-optimizer grad shards into a
+    fresh buffer (grad_data is read, not mutated) and return per-param views."""
+    grad_map: dict[torch.nn.Parameter, torch.Tensor] = {}
+    # Bucket iteration copied from indep_dp.allreduce_grads_and_losses_across_replicas,
+    # which cross-cell all-reduces these same bucket.grad_data buffers.
+    bucket_groups = list(getattr(model_chunk, "bucket_groups", [])) + list(
+        getattr(model_chunk, "expert_parallel_bucket_groups", [])
+    )
+    for bucket_group in bucket_groups:
+        if not bucket_group.ddp_config.use_distributed_optimizer:
+            continue
+        # Same group/size/rank Megatron's grad reduce-scatter uses
+        # (Megatron-LM param_and_grad_buffer.py _ParamAndGradBucketGroup.start_grad_sync).
+        group = bucket_group.intra_distributed_optimizer_instance_group
+        instance_size = bucket_group.intra_distributed_optimizer_instance_size
+        instance_rank = bucket_group.intra_distributed_optimizer_instance_rank
+        for bucket in bucket_group.buckets:
+            grad_data = bucket.grad_data
+            if instance_size > 1:
+                full = torch.empty_like(grad_data)
+                # shard slicing copied from Megatron shard_buffer(); local_shard is
+                # this rank's owned (reduce-scattered) slice.
+                shard_numel = grad_data.numel() // instance_size
+                local_shard = grad_data[instance_rank * shard_numel : (instance_rank + 1) * shard_numel]
+                # all-gather copied from Megatron start_param_sync (it does this on
+                # bucket.param_data); here on grad, into a fresh buffer (grad_data read-only).
+                dist.all_gather_into_tensor(full, local_shard.contiguous(), group=group)
+            else:
+                full = grad_data
+            flat = full.view(-1)
+            # per-param slice copied from Megatron's own bucket.param_data.view(-1)
+            # [start:end].view(shape), using the bucket-local bucket.param_to_index.
+            for param, (start, end) in bucket.param_to_index.items():
+                grad_map[param] = flat[start:end].view(param.shape)
+
+    def get_grad(param: torch.nn.Parameter) -> torch.Tensor | None:
+        reduced = grad_map.get(param)
+        if reduced is not None:
+            return reduced
+        # fallback copied from sglang dumper's original grad read (.grad else main_grad).
+        return param.grad if param.grad is not None else getattr(param, "main_grad", None)
+
+    return get_grad
+
+
 def _log_model_grad_coverage(model: torch.nn.Module) -> None:
     missing: list[str] = []
     with_grad = 0
@@ -174,8 +265,17 @@ def _wrapped(*args: Any, **kwargs: Any) -> Any:
 
 
 def _cleanup_dump_dir(dump_dir: Path) -> None:
-    if _get_rank() == 0 and dump_dir.is_dir():
-        shutil.rmtree(dump_dir)
+    # Best-effort: stale handles (NFS .nfsXXXX stubs) can make rmtree fail with
+    # "Directory not empty"; we don't want that to propagate up and mark the cell
+    # as errored.
+    if (_get_rank() == 0) and dump_dir.is_dir():
+        try:
+            shutil.rmtree(dump_dir)
+        except OSError:
+            logger.warning("dump dir cleanup failed; continuing", exc_info=True)
+
+
+def _barrier_after_dump_dir_cleanup() -> None:
     if dist.is_initialized():
         dist.barrier()
 

tests/e2e/conftest_dumper.py

@@ -195,7 +195,7 @@ def check_dump_dir(
     assert phase_dir.exists(), f"Missing dump dir: {phase_dir}"
     dump_subdirs: list[Path] = list(phase_dir.glob(exp_pattern))
     assert len(dump_subdirs) > 0, f"No {exp_pattern} subdirs in {phase_dir}"
-    dump_files: list[Path] = list(dump_subdirs[0].glob("*.pt"))
+    dump_files: list[Path] = list(dump_subdirs[0].rglob("*.pt"))
     assert len(dump_files) > 0, f"No .pt files in {dump_subdirs[0]}"
     sample: dict = torch.load(dump_files[0], weights_only=False)
     assert isinstance(sample, dict), f"Unexpected type: {type(sample)}"