diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 4535846f637..76e50afe16c 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/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="", ) ) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index a4e91c91c65..dbd3c6e566f 100644 --- a/miles/backends/megatron_utils/model.py +++ b/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. diff --git a/miles/utils/dumper_utils.py b/miles/utils/dumper_utils.py index aea925aa6e3..d0bb59a6e2c 100644 --- a/miles/utils/dumper_utils.py +++ b/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() diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py index dbe4322a849..dfaa1f09981 100644 --- a/tests/e2e/conftest_dumper.py +++ b/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)}"