From b7306b19f6c565716834e14884deb2794db446f0 Mon Sep 17 00:00:00 2001 From: dydy Date: Fri, 21 Aug 2026 17:15:16 +0800 Subject: [PATCH 1/6] feat: full cuda graph --- .../core/fusions/fused_mla_yarn_rope_apply.py | 15 +++++- megatron/core/models/gpt/gpt_layer_specs.py | 7 ++- megatron/core/pipeline_parallel/schedules.py | 17 +++++-- megatron/core/tensor_parallel/layers.py | 8 +++- .../core/transformer/transformer_config.py | 3 ++ .../core/transformer/transformer_layer.py | 47 +++++++++++++++---- 6 files changed, 78 insertions(+), 19 deletions(-) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..0f757ea80b9 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -270,7 +270,14 @@ def forward( cp_rank, cp_size, ) - ctx.save_for_backward(cos, sin) + # Clone cos/sin before saving: the rotary_pos_emb table is shared across + # the Q and KV RoPE autograd functions (same tensor passed to both), and + # an in-place op elsewhere in the captured layer bumps the table's + # autograd version after this save, making backward raise "variable + # needed for gradient computation has been modified by an inplace + # operation". Saving a private clone detaches our saved copy from any + # later in-place mutation of the shared table. + ctx.save_for_backward(cos.clone(), sin.clone()) ctx.qk_head_dim = qk_head_dim ctx.emb_dim = emb_dim ctx.cu_seqlens_q = cu_seqlens_q @@ -658,7 +665,11 @@ def forward( cp_rank, cp_size, ) - ctx.save_for_backward(cos, sin) + # Clone cos/sin before saving -- see the same note in + # ApplyMLARotaryEmbQ.forward. The KV path and the Q path share one + # rotary_pos_emb table; without a clone a later in-place op on the + # shared table breaks this backward too. + ctx.save_for_backward(cos.clone(), sin.clone()) ctx.rotary_interleaved = rotary_interleaved ctx.emb_dim = emb_dim ctx.k_dim = k_dim diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index 984840b3a87..e10be99df86 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -38,6 +38,7 @@ MlpBuilder, TransformerLayer, TransformerLayerSubmodules, + MoETransformerLayer, get_transformer_layer_offset, ) from megatron.core.typed_torch import copy_signature, not_none @@ -640,6 +641,10 @@ def get_gpt_decoder_layer_specs( kitchen_attention_backend=config.kitchen_attention_backend, ) + # MoE layers need the specialized transformer wrapper so local partial + # CUDA Graphs can split router/postprocess from eager dispatch. + moe_layer_spec.module = MoETransformerLayer + # Parse config.moe_layer_freq to determine the pattern of expert/dense layers. # 0 stands for dense layers, 1 stands for expert layers. # For integer N: Creates a pattern with one expert layer every N layers. @@ -767,7 +772,7 @@ def get_gpt_mtp_block_spec_for_backend( if isinstance(spec, TransformerBlockSubmodules): # get the spec for the last layer of decoder block transformer_layer_spec = spec.layer_specs[-1] - elif isinstance(spec, ModuleSpec) and spec.module == TransformerLayer: + elif isinstance(spec, ModuleSpec) and issubclass(spec.module, TransformerLayer): transformer_layer_spec = spec else: raise ValueError(f"Invalid spec: {spec}") diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 29d36ab2c7d..77d33353a72 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -328,17 +328,24 @@ def forward_step_calc_loss( outputs = loss_func(output_tensor) if len(outputs) == 3: output_tensor, num_tokens, loss_reduced = outputs - if not config.calculate_per_token_loss: + if not config.calculate_per_token_loss and not getattr( + output_tensor, '_is_loss_head_final', False + ): # Protect against division by zero when all tokens are masked # in a microbatch. - output_tensor /= torch.clamp(num_tokens, min=1) - output_tensor /= num_microbatches + # NOTE: non-inplace: with the graphed loss head the loss + # scalar is the output of _CudagraphRecordNode (a cudagraph + # buffer view); inplace div on it is rejected by autograd. + # When _is_loss_head_final is set the normalization was + # already computed inside the loss-head cuda graph. + output_tensor = output_tensor / torch.clamp(num_tokens, min=1) + output_tensor = output_tensor / num_microbatches else: # preserve legacy loss averaging behavior (ie, over the number of microbatches) assert len(outputs) == 2 output_tensor, loss_reduced = outputs - output_tensor *= cp_group_size - output_tensor /= num_microbatches + output_tensor = output_tensor * cp_group_size + output_tensor = output_tensor / num_microbatches forward_data_store.append(loss_reduced) else: data = loss_func(output_tensor, non_loss_data=True) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 3ff635b362b..8725021c991 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -349,7 +349,13 @@ def forward(self, input_): weight = GTPEmbeddingWeight.apply(self.weight) # Get the embeddings. - if self.deterministic_mode: + # MUSA: F.embedding's backward (embedding_dense_backward) lowers to a + # thrust sort/unique that SYNCHRONIZES -- illegal inside graph capture + # ("unique_by_key: musaErrorStreamCaptureUnsupported", kills + # cuda_graph_impl=full_iteration; verified by miniexp/unique_bwd_probe). + # The index-based gather (weight[ids]) is numerically identical and its + # backward captures cleanly, so prefer it on MUSA. + if self.deterministic_mode or torch.musa.is_available(): output_parallel = weight[masked_input] else: # F.embedding currently has a non-deterministic backward function diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4f9de546161..11baea5c5a4 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1044,6 +1044,9 @@ class TransformerConfig(ModelParallelConfig): cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when cuda_graph_impl="full_iteration".""" + gpt_postprocess_graph: bool = False + """Capture the GPT output/loss postprocess in a separate local CUDA graph.""" + cuda_graph_modules: Union[str, CudaGraphModule, List[str], List[CudaGraphModule]] = "full" """Selects training capture coverage within per-layer CUDA graphs (local and transformer_engine implementations). diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 3fa91068769..7a41694517c 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1526,6 +1526,12 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + # GraphableMegatronModule invokes this override before TransformerLayer + # has built ``mlp``. Finish MoE-specific graph manager construction only + # after the full layer (including its MoELayer) is initialized. + if self.config.cuda_graph_impl == "local": + self.create_mcore_cudagraph_manager(self.config) + def _should_call_local_cudagraph(self, *args, **kwargs): """ Controls whether the full-layer cudagraph_manager captures the entire forward call @@ -1586,26 +1592,47 @@ def create_mcore_cudagraph_manager(self, config): """ Initializes the CUDA graph manager(s) for the MoE layer. - Unlike the standard layer which typically uses a single manager, this method - can configure multiple graph managers if partial CUDA graphs are enabled via - `cuda_graph_modules`. This allows capturing the static parts of the MoE pass - while leaving the expert computation to execute eagerly. + Partial MoE graphs use dedicated managers for attention, router, and + postprocess. No layer-level manager is installed in that mode, because + it would capture the eager expert dispatch between the MoE phases. """ assert self.config.cuda_graph_impl == "local" from megatron.core.transformer.cuda_graphs import CudaGraphManager + # The base GraphableMegatronModule initializer calls this hook before + # TransformerLayer has constructed the MLP. The subclass calls it again + # after initialization from __init__. + if not hasattr(self, "mlp"): + return + + partial_moe = ( + CudaGraphModule.moe_router in self.config.cuda_graph_modules + or CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules + ) + + if partial_moe: + # The base initializer may have registered an attention manager using + # the generic layer path. It would wrap the whole TransformerLayer, + # so replace it with a manager bound only to _forward_attention. + self._modules.pop("cudagraph_manager", None) + if ( + CudaGraphModule.attn in self.config.cuda_graph_modules + and not hasattr(self, "cudagraph_manager_attn") + ): + self.cudagraph_manager_attn = CudaGraphManager( + config, self, function_name="_forward_attention" + ) + self.transition_cudagraph_scope("partial") + return + if ( not self.config.cuda_graph_modules and self.config.inference_cuda_graph_scope != InferenceCudaGraphScope.block ) or CudaGraphModule.moe in self.config.cuda_graph_modules: - self.cudagraph_manager = CudaGraphManager(config) - elif ( - CudaGraphModule.moe_router in self.config.cuda_graph_modules - or CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules - ): - self.transition_cudagraph_scope('partial') + if not hasattr(self, "cudagraph_manager"): + self.cudagraph_manager = CudaGraphManager(config) def _resolve_token_dispatcher_attr(self, attr_name: str) -> tuple[Any, str]: parent_attr_name, _, leaf_attr_name = attr_name.rpartition('.') From dd5b6ae9a96f03c7f5ee73f491624f8d62824d34 Mon Sep 17 00:00:00 2001 From: dydy Date: Tue, 25 Aug 2026 18:31:15 +0800 Subject: [PATCH 2/6] feat: cpu offload with cuda graph --- .../fine_grained_activation_offload.py | 330 +++++++++++++++++- megatron/core/transformer/cuda_graphs.py | 154 +++++--- .../core/transformer/transformer_config.py | 14 +- megatron/training/training.py | 5 +- 4 files changed, 439 insertions(+), 64 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index bcc8c14ebb6..33894552a87 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import os from collections import defaultdict, deque from contextlib import nullcontext from typing import Any, Dict, Optional, Tuple @@ -397,6 +398,159 @@ def update_offload_info(self, tensor): self.total_tensor_count += 1 +class LocalCudaGraphOffloadGroup: + """Fixed-address VMM slots for one local-graph activation group.""" + + def __init__(self, name: str, min_tensor_size: int): + self.name = name + self.min_tensor_size = min_tensor_size + self.descriptors = [] + self.allocations = [] + self.device_tensors = [] + self.host_tensors = [] + self.capture_index = 0 + self.state = "discovering" + self.d2h_event = torch.cuda.Event() + self.h2d_event = torch.cuda.Event() + + def _eligible(self, tensor): + return ( + torch.is_tensor(tensor) + and not isinstance(tensor, torch.nn.Parameter) + and tensor.device.type == "musa" + and tensor.numel() >= self.min_tensor_size + and not _te_do_not_offload(tensor) + and not getattr(tensor, "_do_not_offload", False) + ) + + @staticmethod + def _descriptor(tensor): + return ( + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) + + def discover_tensor(self, tensor): + """Record an eligible saved-tensor descriptor without changing eager autograd.""" + if self._eligible(tensor): + self.descriptors.append(self._descriptor(tensor)) + return tensor + + @staticmethod + def load_tensor(tensor): + """Return the tensor represented by a saved-tensor hook payload.""" + return tensor + + def allocate(self): + """Allocate all fixed-address slots outside graph capture.""" + from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + + assert self.state == "discovering" + for shape, stride, dtype, device in self.descriptors: + allocation = MUSAActivationVMMAllocation(shape, stride, dtype, device) + self.allocations.append(allocation) + self.device_tensors.append(allocation.tensor) + self.host_tensors.append( + torch.empty_strided(shape, stride, dtype=dtype, device="cpu", pin_memory=True) + ) + self.state = "mapped" + + def begin_capture(self): + """Reset ordered descriptor consumption before forward graph capture.""" + assert self.state == "mapped" + self.capture_index = 0 + + def save_tensor(self, tensor): + """Capture a D2D write and save the matching VMM tensor for backward.""" + if not self._eligible(tensor): + return tensor + if self.capture_index >= len(self.descriptors): + raise RuntimeError(f"{self.name}: capture saved more tensors than discovery") + descriptor = self._descriptor(tensor) + expected = self.descriptors[self.capture_index] + if descriptor != expected: + raise RuntimeError( + f"{self.name}: activation descriptor drift at slot {self.capture_index}: " + f"expected={expected}, actual={descriptor}" + ) + slot_tensor = self.device_tensors[self.capture_index] + slot_tensor.copy_(tensor) + PipelineOffloadManager._local_graph_debug( + f"capture group={self.name} slot={self.capture_index} " + f"source={hex(tensor.data_ptr())} vmm_va={hex(slot_tensor.data_ptr())} " + f"shape={tuple(tensor.shape)} stride={tuple(tensor.stride())} dtype={tensor.dtype} " + f"aligned_bytes={self.allocations[self.capture_index].aligned_bytes}" + ) + self.capture_index += 1 + return slot_tensor + + def finish_capture(self): + """Ensure capture consumed exactly the descriptors discovered during warmup.""" + if self.capture_index != len(self.descriptors): + raise RuntimeError( + f"{self.name}: capture consumed {self.capture_index}/{len(self.descriptors)} slots" + ) + + @property + def logical_bytes(self): + return sum(tensor.numel() * tensor.element_size() for tensor in self.device_tensors) + + @property + def physical_bytes(self): + return sum(allocation.aligned_bytes for allocation in self.allocations) + + def offload_and_release(self, d2h_stream, compute_stream): + """Copy slots to pinned host buffers, then release their physical backing.""" + if not self.device_tensors: + return + if self.state != "mapped": + raise RuntimeError(f"{self.name}: cannot offload slots in state {self.state}") + d2h_stream.wait_stream(compute_stream) + with torch.cuda.stream(d2h_stream): + for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): + host_tensor.copy_(device_tensor, non_blocking=True) + self.d2h_event.record(d2h_stream) + self.d2h_event.synchronize() + before = self._driver_free_memory() + for allocation in self.allocations: + allocation.unmap_and_release() + after = self._driver_free_memory() + self.state = "unmapped" + PipelineOffloadManager._local_graph_debug( + f"release group={self.name} logical_bytes={self.logical_bytes} " + f"physical_bytes={self.physical_bytes} driver_free_delta={after - before}" + ) + + def remap_and_reload(self, h2d_stream, compute_stream): + """Remap fresh backing at the fixed VAs and restore slot contents.""" + if not self.device_tensors: + return + if self.state != "unmapped": + raise RuntimeError(f"{self.name}: cannot reload slots in state {self.state}") + before = self._driver_free_memory() + for allocation in self.allocations: + allocation.create_and_remap() + after = self._driver_free_memory() + with torch.cuda.stream(h2d_stream): + for device_tensor, host_tensor in zip(self.device_tensors, self.host_tensors): + device_tensor.copy_(host_tensor, non_blocking=True) + self.h2d_event.record(h2d_stream) + compute_stream.wait_event(self.h2d_event) + self.state = "mapped" + PipelineOffloadManager._local_graph_debug( + f"remap group={self.name} logical_bytes={self.logical_bytes} " + f"physical_bytes={self.physical_bytes} driver_free_delta={after - before}" + ) + + @staticmethod + def _driver_free_memory(): + from transformer_engine.pytorch.vmm_activation import vmm_driver_memory_info + + return vmm_driver_memory_info()["free_bytes"] + + class PipelineOffloadManager: """ Singleton manager for coordinating activation offloading across pipeline stages. @@ -449,6 +603,17 @@ def __init__(self): # Index of the current forward chunk in the cached chunks forward. self._cached_chunks_index_forward = 0 + # Whole-layer local CUDA graph capture records fixed activation slots here. These are + # separate from eager chunk state because schedule reset runs before graph creation. + self._local_graph_capture_runner = None + self._local_graph_mode = None + self._local_graph_capture_groups = [] + self._local_graph_capture_group = None + self._local_graph_group_index = 0 + self._local_graph_saved_tensors_hooks = None + self._local_graph_d2h_bytes = 0 + self._local_graph_h2d_bytes = 0 + self.do_offload = True # Do not offload the last X groups so that the reloading won't block the computing stream. @@ -489,6 +654,147 @@ def cpu_tensor_pool(self): """Get the shared CPU tensor pool.""" return self._cpu_tensor_pool + @property + def in_local_graph_capture(self): + """Whether a whole-layer local graph is currently recording activation slots.""" + return self._local_graph_capture_runner is not None + + @staticmethod + def _local_graph_debug(message): + if os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_DEBUG") != "1": + return + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + print(f"[local-full-graph-offload] {message}", flush=True) + + @staticmethod + def _local_graph_group_bytes(groups): + return sum(group.logical_bytes for group in groups) + + def begin_local_graph_discovery(self, runner): + """Discover saved activation descriptors during the final eager warmup.""" + assert not self.in_local_graph_capture, "Nested local graph offload discovery" + self._local_graph_capture_runner = runner + self._local_graph_mode = "discovery" + self._local_graph_capture_groups = [] + self._local_graph_capture_group = None + self._local_graph_group_index = 0 + + def end_local_graph_discovery(self, success=True): + """Allocate fixed-address slots after descriptor discovery.""" + runner = self._local_graph_capture_runner + if runner is None or self._local_graph_mode != "discovery": + return + assert self._local_graph_capture_group is None, "Unclosed discovery group" + groups = list(self._local_graph_capture_groups) + if success: + for group in groups: + group.allocate() + runner.local_graph_offload_groups = groups + self._local_graph_debug( + f"discovered groups={len(groups)} logical_bytes={self._local_graph_group_bytes(groups)} " + f"physical_bytes={sum(group.physical_bytes for group in groups)}" + ) + else: + runner.local_graph_offload_groups = [] + self._local_graph_capture_runner = None + self._local_graph_mode = None + self._local_graph_capture_groups = [] + self._local_graph_group_index = 0 + + def begin_local_graph_capture(self, runner): + """Bind graph capture to slots allocated during descriptor discovery.""" + assert not self.in_local_graph_capture, "Nested local CUDA graph offload capture" + groups = getattr(runner, "local_graph_offload_groups", []) + self._local_graph_capture_runner = runner + self._local_graph_mode = "capture" + self._local_graph_capture_groups = groups + self._local_graph_capture_group = None + self._local_graph_group_index = 0 + for group in groups: + group.begin_capture() + + def end_local_graph_capture(self, success=True): + """Validate capture-time descriptor consumption and retain runner slots.""" + runner = self._local_graph_capture_runner + if runner is None or self._local_graph_mode != "capture": + return + assert self._local_graph_capture_group is None, "Unclosed local graph offload group" + groups = self._local_graph_capture_groups + if success: + if self._local_graph_group_index != len(groups): + raise RuntimeError( + f"local graph capture consumed {self._local_graph_group_index}/{len(groups)} groups" + ) + for group in groups: + group.finish_capture() + self._local_graph_debug( + f"captured groups={len(groups)} bytes={self._local_graph_group_bytes(groups)}" + ) + else: + runner.local_graph_offload_groups = [] + self._local_graph_capture_runner = None + self._local_graph_mode = None + self._local_graph_capture_groups = [] + self._local_graph_group_index = 0 + + def local_graph_group_start(self, name, min_tensor_size): + """Open saved-tensor hooks for one discovery or capture group.""" + assert self.in_local_graph_capture + assert self._local_graph_capture_group is None, "Overlapping local graph offload groups" + if self._local_graph_mode == "discovery": + group = LocalCudaGraphOffloadGroup(name, min_tensor_size) + pack_hook = group.discover_tensor + else: + if self._local_graph_group_index >= len(self._local_graph_capture_groups): + raise RuntimeError(f"capture produced unexpected activation group {name}") + group = self._local_graph_capture_groups[self._local_graph_group_index] + if group.name != name or group.min_tensor_size != min_tensor_size: + raise RuntimeError( + f"local graph activation group drift: expected={group.name}, actual={name}" + ) + pack_hook = group.save_tensor + self._local_graph_capture_group = group + self._local_graph_saved_tensors_hooks = saved_tensors_hooks(pack_hook, group.load_tensor) + self._local_graph_saved_tensors_hooks.__enter__() + + def local_graph_group_context_exit(self): + """Close the current group's saved-tensor hooks.""" + if self._local_graph_saved_tensors_hooks is not None: + self._local_graph_saved_tensors_hooks.__exit__(None, None, None) + self._local_graph_saved_tensors_hooks = None + + def local_graph_group_commit(self, name): + """Commit one ordered discovery or capture group.""" + group = self._local_graph_capture_group + assert group is not None and group.name == name + assert self._local_graph_saved_tensors_hooks is None, "Offload context must exit before commit" + if self._local_graph_mode == "discovery": + self._local_graph_capture_groups.append(group) + else: + group.finish_capture() + self._local_graph_group_index += 1 + self._local_graph_capture_group = None + + def local_graph_forward_replay(self, runner): + """D2H-save and physically release VMM slots after forward graph replay.""" + compute_stream = torch.cuda.current_stream() + groups = getattr(runner, "local_graph_offload_groups", []) + for group in groups: + group.offload_and_release(self.d2h_stream, compute_stream) + step_bytes = self._local_graph_group_bytes(groups) + self._local_graph_d2h_bytes += step_bytes + self._local_graph_debug(f"d2h bytes={step_bytes} total_d2h={self._local_graph_d2h_bytes}") + + def local_graph_backward_replay(self, runner): + """Remap and restore VMM slots before backward graph replay.""" + compute_stream = torch.cuda.current_stream() + groups = getattr(runner, "local_graph_offload_groups", []) + for group in reversed(groups): + group.remap_and_reload(self.h2d_stream, compute_stream) + step_bytes = self._local_graph_group_bytes(groups) + self._local_graph_h2d_bytes += step_bytes + self._local_graph_debug(f"h2d bytes={step_bytes} total_h2d={self._local_graph_h2d_bytes}") + def push_offload_groups(self, group_hook, name, forced_released_tensors): """Push the offload groups to the delayed queue.""" debug_rank(f"pushing offload groups to the delayed queue") @@ -954,7 +1260,7 @@ def _can_manage_tensor_for_offload(tensor): return ( not isinstance(tensor, torch.nn.Parameter) and not torch_stray_tensor - and tensor.device.type == "cuda" + and tensor.device.type in ("cuda", "musa") ) def tensor_push(self, tensor): @@ -1047,7 +1353,6 @@ def bulk_reload_group(self): # Only reload if tensor was offloaded (stored as tuple) if isinstance(state, tuple): recovered_tensor = self.reload(state) - debug_rank(f"----recovered_tensor {recovered_tensor.shape}") group_to_reload.push_tensor(tensor_tag, recovered_tensor) group_to_reload.record_reload_event(self.h2d_stream) self._groups_to_reload.pop() @@ -1360,18 +1665,29 @@ def __init__(self, offload: bool, tensor: torch.Tensor, name: str): self.offload = offload self.tensor = tensor self.name = name + self._local_graph_capture = False def __enter__(self): """Enter context manager to enable activation offloading hooks.""" if self.offload: - self.tensor = fine_grained_offloading_group_start(self.tensor, self.name) - PipelineOffloadManager.get_instance().__enter__() + manager = PipelineOffloadManager.get_instance() + if manager.in_local_graph_capture: + self._local_graph_capture = True + min_tensor_size = manager._local_graph_capture_runner.base_module.config.min_offloaded_tensor_size + manager.local_graph_group_start(self.name, min_tensor_size) + else: + self.tensor = fine_grained_offloading_group_start(self.tensor, self.name) + manager.__enter__() return self.tensor def __exit__(self, *args: Any): """Exit context manager to disable activation offloading hooks.""" if self.offload: - PipelineOffloadManager.get_instance().__exit__() + manager = PipelineOffloadManager.get_instance() + if self._local_graph_capture: + manager.local_graph_group_context_exit() + else: + manager.__exit__() @staticmethod def cuda_graph_stream(): @@ -1412,6 +1728,10 @@ def get_context(flag): def group_offload(self, tensor, forced_released_tensors=None, delay_offload=False): """Group offload the tensors.""" if self.offload: + manager = PipelineOffloadManager.get_instance() + if self._local_graph_capture: + manager.local_graph_group_commit(self.name) + return tensor return fine_grained_offloading_group_offload( tensor, self.name, forced_released_tensors, delay_offload ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 782fd6bf14f..d7e33b5d305 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -811,6 +811,8 @@ def forward(ctx, runner, is_first_microbatch, *inputs): torch.cuda.current_stream().wait_event(runner.fwd_completion_event) else: runner.fwd_graph.replay() + if runner.local_graph_offload_groups: + runner._local_graph_offload_manager().local_graph_forward_replay(runner) return runner.fwd_graph_output_surface @staticmethod @@ -835,6 +837,9 @@ def backward(ctx, *grads): if user_output_grad.data_ptr() != cudagraph_output_grad.data_ptr(): cudagraph_output_grad.copy_(user_output_grad) + if runner.local_graph_offload_groups: + runner._local_graph_offload_manager().local_graph_backward_replay(runner) + if runner.use_stream: runner.stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(runner.stream): @@ -902,6 +907,7 @@ def __init__( self.fwd_graph = None self.bwd_graph = None self.bwd_graph_replay_complete_event = torch.cuda.Event() + self.local_graph_offload_groups = [] self.fwd_graph_recorded = False self.bwd_graph_recorded = False @@ -1039,6 +1045,26 @@ def __str__(self): tuple(self.fwd_graph_input_kwarg_metas["hidden_states"].shape), ) + def _uses_local_graph_activation_offload(self): + """Whether this runner captures the whole layer with MoE-only activation offload.""" + config = getattr(self.base_module, "config", None) + return bool( + config is not None + and config.cuda_graph_impl == "local" + and not config.cuda_graph_modules + and config.fine_grained_activation_offloading + and set(config.offload_modules or []) + <= {"expert_fc1", "moe_act", "fused_group_mlp"} + ) + + @staticmethod + def _local_graph_offload_manager(): + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + return PipelineOffloadManager.get_instance() + def get_quantization_context(self): """Return appropriate quantization context (FP8 or FP4) in cudagraph mode.""" if self.fp8_runtime_enabled: @@ -1255,37 +1281,60 @@ def _resolve_input_buffer(ten): ctx = torch.no_grad() if not self.grad_enabled else nullcontext() with ctx: + local_offload_manager = None + local_offload_was_enabled = False + if self._uses_local_graph_activation_offload(): + local_offload_manager = self._local_graph_offload_manager() + local_offload_was_enabled = local_offload_manager.do_offload + local_offload_manager.disable_offload() + # warmup again as case graph capture mode may execute a different codepath _set_warmup_start() - for _ in range(self.num_warmup_steps): - with self.get_quantization_context(): - - def clone_ten(ten): - if not torch.is_tensor(ten): - return ten - return torch.clone(ten).detach().requires_grad_(ten.requires_grad) - - warmup_args = tree_map(clone_ten, self.fwd_graph_input_args) - warmup_kwargs = tree_map(clone_ten, self.fwd_graph_input_kwargs) - warmup_outputs = self.func(*warmup_args, **warmup_kwargs) - - if self.grad_enabled: - warmup_outputs = self.get_tensors(warmup_outputs) - warmup_outputs = tuple(o for o in warmup_outputs if o.requires_grad) - input_tensors = self.get_tensors(warmup_args, warmup_kwargs) - torch.autograd.grad( - outputs=warmup_outputs, - inputs=tuple(i for i in input_tensors if i.requires_grad), - grad_outputs=tuple(torch.zeros_like(o) for o in warmup_outputs), - only_inputs=True, - allow_unused=True, - ) + for warmup_index in range(self.num_warmup_steps): + discovering_local_slots = ( + local_offload_manager is not None + and warmup_index == self.num_warmup_steps - 1 + ) + if discovering_local_slots: + local_offload_manager.begin_local_graph_discovery(self) + try: + with self.get_quantization_context(): + + def clone_ten(ten): + if not torch.is_tensor(ten): + return ten + return torch.clone(ten).detach().requires_grad_(ten.requires_grad) + + warmup_args = tree_map(clone_ten, self.fwd_graph_input_args) + warmup_kwargs = tree_map(clone_ten, self.fwd_graph_input_kwargs) + warmup_outputs = self.func(*warmup_args, **warmup_kwargs) + + if self.grad_enabled: + warmup_outputs = self.get_tensors(warmup_outputs) + warmup_outputs = tuple(o for o in warmup_outputs if o.requires_grad) + input_tensors = self.get_tensors(warmup_args, warmup_kwargs) + torch.autograd.grad( + outputs=warmup_outputs, + inputs=tuple(i for i in input_tensors if i.requires_grad), + grad_outputs=tuple(torch.zeros_like(o) for o in warmup_outputs), + only_inputs=True, + allow_unused=True, + ) - if self.gtp_remat: - wait_async_comms(GTPChain.GRAPHED.value) - self._sync_against_side_streams(self.bwd_side_streams) + if self.gtp_remat: + wait_async_comms(GTPChain.GRAPHED.value) + self._sync_against_side_streams(self.bwd_side_streams) + except Exception: + if discovering_local_slots: + local_offload_manager.end_local_graph_discovery(success=False) + raise + else: + if discovering_local_slots: + local_offload_manager.end_local_graph_discovery() _set_warmup_end() + if local_offload_manager is not None and local_offload_was_enabled: + local_offload_manager.enable_offload() with self.get_quantization_context(): torch.cuda.synchronize() @@ -1302,30 +1351,42 @@ def clone_ten(ten): if FREEZE_GC: gc.freeze() - with torch.cuda.graph( - self.fwd_graph, pool=self.mempool, capture_error_mode="thread_local" - ): + if local_offload_manager is not None: + local_offload_manager.begin_local_graph_capture(self) - self._sync_against_side_streams(self.fwd_side_streams) + try: + with torch.cuda.graph( + self.fwd_graph, pool=self.mempool, capture_error_mode="thread_local" + ): - fwd_graph_outputs = self.func( - *self.fwd_graph_input_args, **self.fwd_graph_input_kwargs - ) + self._sync_against_side_streams(self.fwd_side_streams) - if self.gtp_remat: - # Forward only issues AG prefetches (no wgrad RS), so drain AG and skip RS. - wait_async_comms(GTPChain.GRAPHED.value, skip_rs=True) + fwd_graph_outputs = self.func( + *self.fwd_graph_input_args, **self.fwd_graph_input_kwargs + ) - if self.fwd_side_streams: - self._wait_side_streams(self.fwd_side_streams) + if self.gtp_remat: + # Forward only issues AG prefetches (no wgrad RS), so drain AG and skip RS. + wait_async_comms(GTPChain.GRAPHED.value, skip_rs=True) - if self.use_stream: - self.fwd_completion_event.record() + if self.fwd_side_streams: + self._wait_side_streams(self.fwd_side_streams) - # Unfreeze GC. - if FREEZE_GC: - gc.unfreeze() + if self.use_stream: + self.fwd_completion_event.record() + except Exception: + if local_offload_manager is not None: + local_offload_manager.end_local_graph_capture(success=False) + raise + else: + if local_offload_manager is not None: + local_offload_manager.end_local_graph_capture() + finally: + # A failed graph capture must not leave GC frozen. + if FREEZE_GC: + gc.unfreeze() + if FREEZE_GC: # gc.collect() drops references to unreachable tensors created during capture, # returning their storage to the allocator to avoid a slowdown during replay. # However, it forces expensive global garbage collection, so must be done @@ -2641,12 +2702,11 @@ def _get_fp8_enabled(): FineGrainedActivationOffloadingInterface as off_interface, ) - # TE CUDA graph warmup should establish graph state without launching - # activation D2H copies; the post-warmup hook restores offloading for - # the measured/replay iterations. + # Keep activation offload disabled during both TE warmup and graph capture. + # The MUSA offload path uses a separate D2H stream, which cannot be left + # running when the graph capture stream is finalized. if self.config.fine_grained_activation_offloading: kwargs['pre_warmup_hook'] = off_interface.disable_offload - kwargs['post_warmup_hook'] = off_interface.enable_offload return kwargs kwargs = get_make_graphed_callables_kwargs() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 11baea5c5a4..1dd19ec043d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2607,25 +2607,19 @@ def _scope_to_str(s): "'fused_group_mlp'. " f"Unsupported offload_modules: {sorted(unsupported_offload_modules)}." ) - assert self.cuda_graph_modules, ( - "fine-grained activation offloading with cuda_graph_impl='local' " - "is not supported with whole-layer CUDA graph capture." - ) - local_partial_moe_offload = ( + local_moe_offload = ( self.cuda_graph_impl == "local" and bool(offload_modules) and offload_modules <= {"expert_fc1", "moe_act", "fused_group_mlp"} - and CudaGraphModule.moe not in self.cuda_graph_modules ) assert ( self.cuda_graph_impl in ("transformer_engine", "full_iteration") - or local_partial_moe_offload + or local_moe_offload ), ( "fine-grained activation offloading is only supported with " "transformer_engine CUDA graph implementation or local CUDA graph " - "implementation with partial MoE offload. Local partial CUDA graphs " - "are supported only for expert_fc1, moe_act, or fused_group_mlp " - "offload when the full MoE module is not captured." + "implementation with MoE expert offload (expert_fc1, moe_act, or " + "fused_group_mlp)." ) assert ( CudaGraphModule.moe not in self.cuda_graph_modules diff --git a/megatron/training/training.py b/megatron/training/training.py index 57b8efe9390..cc7143848e5 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1107,8 +1107,9 @@ def pretrain( timers = get_timers() if args.fine_grained_activation_offloading: - from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu - set_ideal_affinity_for_current_gpu() + if os.getenv("ACCELERATOR_BACKEND", "").lower() != "musa": + from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu + set_ideal_affinity_for_current_gpu() if cfg_container.logger.log_progress: From e9b319b454ceb5b94a53979dd72a7c16b3068317 Mon Sep 17 00:00:00 2001 From: dydy Date: Tue, 1 Sep 2026 10:08:37 +0800 Subject: [PATCH 3/6] feat: ping-pong slot to fast slot offload --- .../fine_grained_activation_offload.py | 394 ++++++++++++++++-- megatron/core/transformer/cuda_graphs.py | 85 +++- ...test_fine_grained_activation_offloading.py | 285 +++++++++++++ .../transformer/test_cuda_graphs.py | 53 +++ 4 files changed, 772 insertions(+), 45 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 33894552a87..aa5d1cecba0 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -15,6 +15,17 @@ from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.utils import nvtx_range_pop, nvtx_range_push +try: + from transformer_engine.pytorch.vmm_activation import ( + release_hooks_after, + remap_and_copy_after, + wait_remap_copy_on_stream, + ) +except ImportError: + release_hooks_after = None # eager-only fallback; local graph paths require TE + remap_and_copy_after = None + wait_remap_copy_on_stream = None + def debug_rank(message): """Print debug message for a specific rank when DEBUG is enabled.""" @@ -412,6 +423,11 @@ def __init__(self, name: str, min_tensor_size: int): self.state = "discovering" self.d2h_event = torch.cuda.Event() self.h2d_event = torch.cuda.Event() + # Async release/remap contexts returned by the batched host functions. + self.release_context = None + self.remap_context = None + self.slot_bank = None + self.resident = False def _eligible(self, tensor): return ( @@ -443,15 +459,30 @@ def load_tensor(tensor): """Return the tensor represented by a saved-tensor hook payload.""" return tensor - def allocate(self): - """Allocate all fixed-address slots outside graph capture.""" + def allocate(self, manager=None, slot_bank=None, group_index=None): + """Allocate fixed-address slots, optionally from a resident shared bank.""" from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation assert self.state == "discovering" - for shape, stride, dtype, device in self.descriptors: - allocation = MUSAActivationVMMAllocation(shape, stride, dtype, device) + self.slot_bank = slot_bank + self.resident = slot_bank is not None + for tensor_index, (shape, stride, dtype, device) in enumerate(self.descriptors): + if self.resident: + allocation = manager.local_graph_resident_allocation( + slot_bank, + self.name, + group_index, + tensor_index, + shape, + stride, + dtype, + device, + ) + else: + allocation = MUSAActivationVMMAllocation(shape, stride, dtype, device) self.allocations.append(allocation) self.device_tensors.append(allocation.tensor) + # Host storage remains runner-local even when the GPU slot is shared. self.host_tensors.append( torch.empty_strided(shape, stride, dtype=dtype, device="cpu", pin_memory=True) ) @@ -501,8 +532,73 @@ def logical_bytes(self): def physical_bytes(self): return sum(allocation.aligned_bytes for allocation in self.allocations) - def offload_and_release(self, d2h_stream, compute_stream): - """Copy slots to pinned host buffers, then release their physical backing.""" + def rebind_resident(self, manager, slot_bank, group_index): + """Replace provisional capture allocations with a shared resident bank.""" + if self.resident: + return + from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + + old_allocations = self.allocations + self.allocations = [] + self.device_tensors = [] + self.resident = True + self.slot_bank = slot_bank + for tensor_index, (shape, stride, dtype, device) in enumerate(self.descriptors): + allocation = manager.local_graph_resident_allocation( + slot_bank, self.name, group_index, tensor_index, shape, stride, dtype, device + ) + self.allocations.append(allocation) + self.device_tensors.append(allocation.tensor) + for allocation in old_allocations: + allocation.close() + + def capture_snapshot(self): + """Preserve capture-time values before another runner reuses this bank.""" + if not self.resident: + return + for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): + host_tensor.copy_(device_tensor, non_blocking=False) + + def capture_restore(self): + """Restore capture-time values before capturing the backward graph.""" + if not self.resident: + return + for device_tensor, host_tensor in zip(self.device_tensors, self.host_tensors): + device_tensor.copy_(host_tensor, non_blocking=False) + + def enqueue_resident_d2h(self, d2h_stream, compute_stream): + """Save a resident slot without releasing its fixed mapping.""" + if self.state != "mapped": + raise RuntimeError(f"{self.name}: cannot save resident slots in state {self.state}") + d2h_stream.wait_stream(compute_stream) + with torch.cuda.stream(d2h_stream): + for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): + host_tensor.copy_(device_tensor, non_blocking=True) + self.state = "resident_d2h_pending" + + def enqueue_resident_h2d(self, h2d_stream): + """Restore host values into an always-mapped resident slot.""" + if self.state not in ("resident_d2h_pending", "mapped"): + raise RuntimeError(f"{self.name}: cannot reload resident slots in state {self.state}") + with torch.cuda.stream(h2d_stream): + for device_tensor, host_tensor in zip(self.device_tensors, self.host_tensors): + device_tensor.copy_(host_tensor, non_blocking=True) + self.state = "resident_reload_pending" + + def adopt_resident_reload(self): + """Mark a resident reload consumable after its stream wait is installed.""" + if self.state != "resident_reload_pending": + raise RuntimeError(f"{self.name}: resident reload is in state {self.state}") + self.state = "mapped" + + def enqueue_d2h(self, d2h_stream, compute_stream): + """Queue this group's D2H copies and hand off release to a stream host func. + + A single host function follows all copies on the D2H stream: once the + burst completes it moves the physical unmap/release to a resident + worker thread (driver calls are not allowed on the host-func callback + thread). No event wait or polling is needed afterwards. + """ if not self.device_tensors: return if self.state != "mapped": @@ -512,36 +608,80 @@ def offload_and_release(self, d2h_stream, compute_stream): for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): host_tensor.copy_(device_tensor, non_blocking=True) self.d2h_event.record(d2h_stream) - self.d2h_event.synchronize() - before = self._driver_free_memory() + self.release_context = release_hooks_after(self.allocations, d2h_stream) + self.state = "d2h_pending" + + def try_release(self): + """Adopt the worker's release if it already finished; never blocks.""" + if self.state != "d2h_pending": + return False + if not self._release_context_done(): + return False + return self._adopt_released() + + def wait_and_release(self): + """Wait for the async release then adopt the unmapped state; backward time.""" + # Already released by a forward-time drain: nothing to wait for. + if self.state == "unmapped": + return False + if self.state != "d2h_pending": + raise RuntimeError(f"{self.name}: cannot release slots in state {self.state}") + self._wait_release_context() + return self._adopt_released() + + def _release_context_done(self): + """Whether the worker finished the async release (host-side check).""" + if self.release_context is None: + raise RuntimeError(f"{self.name}: missing async release context in state {self.state}") + return dict(self.release_context.status())["done"] == 1 + + def _wait_release_context(self): + """Block until the worker finished; surface any worker error.""" for allocation in self.allocations: - allocation.unmap_and_release() - after = self._driver_free_memory() + allocation.wait_for_async_release() + + def _adopt_released(self): + """Apply the worker-completed release to this group's bookkeeping.""" + status = dict(self.release_context.status()) + self.release_context = None + if status["error"]: + raise RuntimeError( + f"{self.name}: async release failed on worker: MUresult {status['error_code']}" + ) self.state = "unmapped" PipelineOffloadManager._local_graph_debug( f"release group={self.name} logical_bytes={self.logical_bytes} " - f"physical_bytes={self.physical_bytes} driver_free_delta={after - before}" + f"physical_bytes={self.physical_bytes}" ) + return True - def remap_and_reload(self, h2d_stream, compute_stream): - """Remap fresh backing at the fixed VAs and restore slot contents.""" + def prepare_remap(self, remap_context): + """Bind this group to its runner-level asynchronous remap batch.""" if not self.device_tensors: - return - if self.state != "unmapped": - raise RuntimeError(f"{self.name}: cannot reload slots in state {self.state}") - before = self._driver_free_memory() + return False + if self.state not in ("d2h_pending", "unmapped"): + raise RuntimeError(f"{self.name}: cannot prepare reload in state {self.state}") + self.remap_context = remap_context + self.state = "remap_pending" + return True + + def adopt_reload_submission(self): + """Adopt remaps after the runner installed its GPU-side H2D wait.""" + if self.state != "remap_pending": + raise RuntimeError(f"{self.name}: cannot adopt reload in state {self.state}") for allocation in self.allocations: - allocation.create_and_remap() - after = self._driver_free_memory() - with torch.cuda.stream(h2d_stream): - for device_tensor, host_tensor in zip(self.device_tensors, self.host_tensors): - device_tensor.copy_(host_tensor, non_blocking=True) - self.h2d_event.record(h2d_stream) - compute_stream.wait_event(self.h2d_event) + allocation.adopt_async_remap() + status = dict(self.remap_context.status()) + self.remap_context = None + self.release_context = None + if status["error"]: + raise RuntimeError( + f"{self.name}: async reload failed on worker: error {status['error_code']}" + ) self.state = "mapped" PipelineOffloadManager._local_graph_debug( - f"remap group={self.name} logical_bytes={self.logical_bytes} " - f"physical_bytes={self.physical_bytes} driver_free_delta={after - before}" + f"reload submitted group={self.name} logical_bytes={self.logical_bytes} " + f"physical_bytes={self.physical_bytes}" ) @staticmethod @@ -613,6 +753,10 @@ def __init__(self): self._local_graph_saved_tensors_hooks = None self._local_graph_d2h_bytes = 0 self._local_graph_h2d_bytes = 0 + self._local_graph_resident_slots = [{}, {}] + self._local_graph_bank_d2h_events = [torch.cuda.Event(), torch.cuda.Event()] + self._local_graph_bank_d2h_recorded = [False, False] + self._local_graph_bank_consumed_events = [None, None] self.do_offload = True @@ -670,6 +814,49 @@ def _local_graph_debug(message): def _local_graph_group_bytes(groups): return sum(group.logical_bytes for group in groups) + @staticmethod + def _local_graph_ping_pong_enabled(runner=None): + """Whether this runner uses resident slots for local Graph CPU offload.""" + if runner is None: + return False + enabled = getattr(runner, "local_graph_slot_bank", None) is not None + if enabled and not getattr(runner, "local_graph_offload_groups", []): + # During discovery, groups are not attached until allocation completes. + enabled = True + return enabled + + def prepare_local_graph_ping_pong(self, runner, slot_bank): + """Rebind provisional runner slots to the selected resident bank.""" + runner.local_graph_slot_bank = slot_bank + for group_index, group in enumerate(runner.local_graph_offload_groups): + group.rebind_resident(self, slot_bank, group_index) + + def local_graph_resident_allocation( + self, bank, name, group_index, tensor_index, shape, stride, dtype, device + ): + """Get or create one descriptor-compatible allocation in a resident bank.""" + from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + + key = (name, group_index, tensor_index, tuple(shape), tuple(stride), dtype, device) + allocations = self._local_graph_resident_slots[bank] + if key not in allocations: + allocations[key] = MUSAActivationVMMAllocation(shape, stride, dtype, device) + return allocations[key] + + def local_graph_capture_snapshot(self, runner): + """Preserve shared-slot values needed by this runner's later bwd capture.""" + if not self._local_graph_ping_pong_enabled(runner): + return + for group in runner.local_graph_offload_groups: + group.capture_snapshot() + + def local_graph_capture_restore(self, runner): + """Restore shared-slot values before this runner's backward capture.""" + if not self._local_graph_ping_pong_enabled(runner): + return + for group in runner.local_graph_offload_groups: + group.capture_restore() + def begin_local_graph_discovery(self, runner): """Discover saved activation descriptors during the final eager warmup.""" assert not self.in_local_graph_capture, "Nested local graph offload discovery" @@ -687,8 +874,13 @@ def end_local_graph_discovery(self, success=True): assert self._local_graph_capture_group is None, "Unclosed discovery group" groups = list(self._local_graph_capture_groups) if success: - for group in groups: - group.allocate() + slot_bank = ( + runner.local_graph_slot_bank + if self._local_graph_ping_pong_enabled(runner) + else None + ) + for group_index, group in enumerate(groups): + group.allocate(self, slot_bank, group_index) runner.local_graph_offload_groups = groups self._local_graph_debug( f"discovered groups={len(groups)} logical_bytes={self._local_graph_group_bytes(groups)} " @@ -775,25 +967,153 @@ def local_graph_group_commit(self, name): self._local_graph_group_index += 1 self._local_graph_capture_group = None + def local_graph_forward_wait_ready(self, runner, replay_stream): + """Fence reuse of a resident bank until its preceding D2H finishes.""" + if not self._local_graph_ping_pong_enabled(runner): + return False + bank = runner.local_graph_slot_bank + if self._local_graph_bank_d2h_recorded[bank]: + replay_stream.wait_event(self._local_graph_bank_d2h_events[bank]) + consumed_event = self._local_graph_bank_consumed_events[bank] + if consumed_event is not None: + replay_stream.wait_event(consumed_event) + return True + def local_graph_forward_replay(self, runner): - """D2H-save and physically release VMM slots after forward graph replay.""" + """Enqueue all groups' D2H after forward graph replay.""" compute_stream = torch.cuda.current_stream() groups = getattr(runner, "local_graph_offload_groups", []) - for group in groups: - group.offload_and_release(self.d2h_stream, compute_stream) + active_groups = [group for group in groups if group.device_tensors] step_bytes = self._local_graph_group_bytes(groups) self._local_graph_d2h_bytes += step_bytes + if not active_groups: + self._local_graph_debug(f"d2h bytes={step_bytes} total_d2h={self._local_graph_d2h_bytes}") + return + if self._local_graph_ping_pong_enabled(runner): + bank = runner.local_graph_slot_bank + for group in active_groups: + group.enqueue_resident_d2h(self.d2h_stream, compute_stream) + with torch.cuda.stream(self.d2h_stream): + self._local_graph_bank_d2h_events[bank].record(self.d2h_stream) + self._local_graph_bank_d2h_recorded[bank] = True + else: + # Submit every group's D2H back-to-back on the shared stream, then + # return without a host-side wait. + for group in active_groups: + group.enqueue_d2h(self.d2h_stream, compute_stream) + self.drain_pending_d2h(runner) self._local_graph_debug(f"d2h bytes={step_bytes} total_d2h={self._local_graph_d2h_bytes}") - def local_graph_backward_replay(self, runner): - """Remap and restore VMM slots before backward graph replay.""" - compute_stream = torch.cuda.current_stream() + def drain_pending_d2h(self, runner): + """Release VMM backing for groups whose D2H already completed; never blocks.""" groups = getattr(runner, "local_graph_offload_groups", []) - for group in reversed(groups): - group.remap_and_reload(self.h2d_stream, compute_stream) - step_bytes = self._local_graph_group_bytes(groups) + released = sum(1 for group in groups if group.try_release()) + if released: + self._local_graph_debug(f"lazy released groups={released}") + + @staticmethod + def _local_graph_active_groups(runner): + groups = getattr(runner, "local_graph_offload_groups", []) + return [group for group in reversed(groups) if group.device_tensors] + + def local_graph_backward_prepare(self, runner): + """Submit one runner's resident reload or VMM remap/H2D batch.""" + active_groups = self._local_graph_active_groups(runner) + if not active_groups: + return False + state = getattr(runner, "local_graph_reload_state", None) + if state == "reload_pending": + return True + if state is not None: + raise RuntimeError(f"cannot prepare local graph runner in reload state {state}") + if self._local_graph_ping_pong_enabled(runner): + bank = runner.local_graph_slot_bank + # Before the first backward use, do not overwrite this bank until + # its final forward D2H has saved the activation currently in it. + if self._local_graph_bank_d2h_recorded[bank]: + self.h2d_stream.wait_event(self._local_graph_bank_d2h_events[bank]) + # Later uses wait until the prior backward graph reading this bank + # has completed. Adjacent runners use the opposite bank and overlap. + consumed_event = self._local_graph_bank_consumed_events[bank] + if consumed_event is not None: + self.h2d_stream.wait_event(consumed_event) + for group in active_groups: + group.enqueue_resident_h2d(self.h2d_stream) + reload_event = torch.cuda.Event() + with torch.cuda.stream(self.h2d_stream): + reload_event.record(self.h2d_stream) + runner.local_graph_reload_event = reload_event + runner.local_graph_reload_context = None + else: + # Submit in backward consumption order. The last forward-offloaded + # group has the latest release dependency, so waiting on it first + # makes every earlier dependency a fast path. + allocations = [ + allocation for group in active_groups for allocation in group.allocations + ] + host_tensors = [tensor for group in active_groups for tensor in group.host_tensors] + reload_context = remap_and_copy_after( + allocations, host_tensors, self.h2d_stream + ) + for group in active_groups: + group.prepare_remap(reload_context) + runner.local_graph_reload_context = reload_context + runner.local_graph_reload_state = "reload_pending" + return True + + def local_graph_backward_finish_prepare(self, runner): + """Compatibility no-op: prepare already handed the full reload to the worker.""" + active_groups = self._local_graph_active_groups(runner) + if not active_groups: + return False + if getattr(runner, "local_graph_reload_state", None) != "reload_pending": + raise RuntimeError("local graph reload was not prepared before finish") + return True + + def local_graph_backward_wait_ready(self, runner, compute_stream=None): + """Install the reload event dependency when this runner is consumed.""" + active_groups = self._local_graph_active_groups(runner) + step_bytes = self._local_graph_group_bytes( + getattr(runner, "local_graph_offload_groups", []) + ) self._local_graph_h2d_bytes += step_bytes + if not active_groups: + self._local_graph_debug(f"h2d bytes={step_bytes} total_h2d={self._local_graph_h2d_bytes}") + return False + if getattr(runner, "local_graph_reload_state", None) is None: + self.local_graph_backward_prepare(runner) + if getattr(runner, "local_graph_reload_state", None) != "reload_pending": + raise RuntimeError("local graph runner reload is not pending for backward consumption") + if compute_stream is None: + compute_stream = torch.cuda.current_stream() + if self._local_graph_ping_pong_enabled(runner): + reload_event = runner.local_graph_reload_event + if reload_event is None: + raise RuntimeError("local graph resident reload has no completion event") + compute_stream.wait_event(reload_event) + for group in active_groups: + group.adopt_resident_reload() + runner.local_graph_reload_event = None + else: + reload_context = runner.local_graph_reload_context + wait_remap_copy_on_stream(reload_context, compute_stream) + for group in active_groups: + group.adopt_reload_submission() + runner.local_graph_reload_context = None + runner.local_graph_reload_state = None self._local_graph_debug(f"h2d bytes={step_bytes} total_h2d={self._local_graph_h2d_bytes}") + return True + + def local_graph_backward_mark_consumed(self, runner, consumed_event): + """Publish when a resident bank may be overwritten by a later reload.""" + if not self._local_graph_ping_pong_enabled(runner): + return False + self._local_graph_bank_consumed_events[runner.local_graph_slot_bank] = consumed_event + return True + + def local_graph_backward_replay(self, runner): + """Compatibility fallback for callers without look-ahead scheduling.""" + self.local_graph_backward_wait_ready(runner) def push_offload_groups(self, group_hook, name, forced_released_tensors): """Push the offload groups to the delayed queue.""" diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index d7e33b5d305..6877ac81bab 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -489,6 +489,7 @@ class _CudagraphGlobalRecord: cudagraph_record: list[tuple] = [] cudagraph_inference_record: list[tuple] = [] _saved_tensors_observer = None + backward_runners: list = [] @classmethod def _enable_saved_tensors_observer(cls): @@ -529,6 +530,24 @@ def record_bwd_graph(cls, runner): """Record a bwd graph to 'cudagraph_record""" cls.cudagraph_record.append((runner, "bwd")) + @classmethod + def _link_backward_runners(cls): + """Link runners in the backward order observed by autograd and the schedule.""" + cls.backward_runners = [g[0] for g in cls.cudagraph_record if g[1] == "bwd"] + local_graph_offload_index = 0 + for runner in cls.backward_runners: + runner.next_bwd_runner = None + uses_local_offload = getattr( + runner, "_uses_local_graph_activation_offload", lambda: False + )() + runner.local_graph_slot_bank = ( + local_graph_offload_index % 2 if uses_local_offload else None + ) + if uses_local_offload: + local_graph_offload_index += 1 + for runner, next_runner in zip(cls.backward_runners, cls.backward_runners[1:]): + runner.next_bwd_runner = next_runner + @classmethod def create_cudagraphs(cls): """Create recorded CUDA graphs, then remove the saved-tensor observer.""" @@ -562,6 +581,17 @@ def _create_cudagraphs(cls): [isinstance(m, TransformerEngineBaseModule) for m in base_module.modules()] ) + # Persist the actual autograd/schedule backward order before the global + # creation record is cleared. Local activation reload uses only the + # immediate successor, bounding look-ahead physical backing to one runner. + cls._link_backward_runners() + + for runner in cls.backward_runners: + if runner.local_graph_offload_groups: + runner._local_graph_offload_manager().prepare_local_graph_ping_pong( + runner, runner.local_graph_slot_bank + ) + progress_bar = enumerate(cls.cudagraph_record) time_start = time.time() mem_stats_start = torch.cuda.memory_stats() @@ -698,10 +728,15 @@ def delete_cuda_graphs(): runner.bwd_graph = None runner.mempool = None + for runner in _CudagraphGlobalRecord.backward_runners: + runner.next_bwd_runner = None + runner.local_graph_reload_state = None + # Reset global tracking state _CudagraphGlobalRecord.cudagraph_created = False _CudagraphGlobalRecord.cudagraph_record = [] _CudagraphGlobalRecord.cudagraph_inference_record = [] + _CudagraphGlobalRecord.backward_runners = [] _GTP_RUNNER_STREAMS.clear() # TODO: Optional?: Force garbage collection to clean up memory @@ -804,6 +839,10 @@ def forward(ctx, runner, is_first_microbatch, *inputs): _set_skip_fp8_weight_update_tensor(not is_first_microbatch) runner.fp8_param_cache_updated = is_first_microbatch + if runner.local_graph_offload_groups: + offload_manager = runner._local_graph_offload_manager() + replay_stream = runner.stream if runner.use_stream else torch.cuda.current_stream() + offload_manager.local_graph_forward_wait_ready(runner, replay_stream) if runner.use_stream: runner.stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(runner.stream): @@ -812,7 +851,7 @@ def forward(ctx, runner, is_first_microbatch, *inputs): else: runner.fwd_graph.replay() if runner.local_graph_offload_groups: - runner._local_graph_offload_manager().local_graph_forward_replay(runner) + offload_manager.local_graph_forward_replay(runner) return runner.fwd_graph_output_surface @staticmethod @@ -837,18 +876,38 @@ def backward(ctx, *grads): if user_output_grad.data_ptr() != cudagraph_output_grad.data_ptr(): cudagraph_output_grad.copy_(user_output_grad) + offload_manager = None + next_runner = runner.next_bwd_runner + next_prepared = False if runner.local_graph_offload_groups: - runner._local_graph_offload_manager().local_graph_backward_replay(runner) + offload_manager = runner._local_graph_offload_manager() + replay_stream = runner.stream if runner.use_stream else torch.cuda.current_stream() + offload_manager.local_graph_backward_wait_ready(runner, replay_stream) + if next_runner is not None and next_runner.local_graph_offload_groups: + if offload_manager is None: + offload_manager = next_runner._local_graph_offload_manager() + next_prepared = offload_manager.local_graph_backward_prepare(next_runner) - if runner.use_stream: - runner.stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(runner.stream): + try: + if runner.use_stream: + runner.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(runner.stream): + runner.bwd_graph.replay() + torch.cuda.current_stream().wait_event(runner.bwd_completion_event) + else: runner.bwd_graph.replay() - torch.cuda.current_stream().wait_event(runner.bwd_completion_event) - else: - runner.bwd_graph.replay() + finally: + # The current graph has already been submitted. Waiting for the + # successor's remap worker and enqueueing H2D now overlaps both with + # current-runner GPU execution. + if next_prepared: + offload_manager.local_graph_backward_finish_prepare(next_runner) runner.bwd_graph_replay_complete_event.record(torch.cuda.current_stream()) + if runner.local_graph_offload_groups: + offload_manager.local_graph_backward_mark_consumed( + runner, runner.bwd_graph_replay_complete_event + ) for param in runner.params_to_backprop: param._cudagraph_wgrad_ready_event = runner.bwd_graph_replay_complete_event runner.status = _GraphStatus.FWD_READY @@ -908,6 +967,10 @@ def __init__( self.bwd_graph = None self.bwd_graph_replay_complete_event = torch.cuda.Event() self.local_graph_offload_groups = [] + self.next_bwd_runner = None + self.local_graph_slot_bank = None + self.local_graph_reload_state = None + self.local_graph_reload_event = None self.fwd_graph_recorded = False self.bwd_graph_recorded = False @@ -1381,6 +1444,7 @@ def clone_ten(ten): else: if local_offload_manager is not None: local_offload_manager.end_local_graph_capture() + local_offload_manager.local_graph_capture_snapshot(self) finally: # A failed graph capture must not leave GC frozen. if FREEZE_GC: @@ -1482,6 +1546,11 @@ def create_bwd_graph(self): out_grad = alloc_tensor_from_graph_mempool(o) self.static_grad_outputs.append(out_grad) + # Shared ping-pong slots may have been overwritten while later forward + # graphs were captured. Restore this runner before capturing its consumers. + if self.local_graph_offload_groups: + self._local_graph_offload_manager().local_graph_capture_restore(self) + # Freeze GC, to speed up capture time ~15-20x. if FREEZE_GC: gc.freeze() diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index dce35bdd4fa..bdf0aa1b3b3 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -968,3 +968,288 @@ def test_fine_grained_activation_offloading_with_cuda_graph( finally: Utils.destroy_model_parallel() + + +# ============================================================================= +# Local Graph Lazy-Release D2H Offload Tests (CPU-only, mocked streams/events) +# ============================================================================= + + +class _RecordingEvent: + def __init__(self, log, name): + self._log = log + self.name = name + + def record(self, stream): + self._log.append((f"{self.name}.record", stream)) + + def synchronize(self): + self._log.append((f"{self.name}.synchronize",)) + + def query(self): + self._log.append((f"{self.name}.query",)) + return False + + +class _RecordingStream: + def __init__(self, name, log): + self.name = name + self._log = log + + def wait_stream(self, other): + self._log.append((f"{self.name}.wait_stream", other)) + + def wait_event(self, event): + self._log.append((f"{self.name}.wait_event", event)) + + +class _RecordingGraphGroup: + """Stand-in for LocalCudaGraphOffloadGroup in replay orchestration tests.""" + + def __init__(self, name, log, num_tensors=1, tensor_bytes=1024): + self.name = name + self._log = log + self.device_tensors = [object()] * num_tensors + self.allocations = [(name, "allocation", index) for index in range(num_tensors)] + self.host_tensors = [(name, "host", index) for index in range(num_tensors)] + self._logical_bytes = num_tensors * tensor_bytes + self.state = "mapped" + + @property + def logical_bytes(self): + return self._logical_bytes + + def enqueue_d2h(self, d2h_stream, compute_stream): + self._log.append(("group.enqueue_d2h", self.name)) + + def enqueue_resident_d2h(self, d2h_stream, compute_stream): + self._log.append(("group.enqueue_resident_d2h", self.name, d2h_stream.name, compute_stream)) + self.state = "resident_d2h_pending" + + def enqueue_resident_h2d(self, h2d_stream): + self._log.append(("group.enqueue_resident_h2d", self.name, h2d_stream.name)) + self.state = "resident_reload_pending" + + def adopt_resident_reload(self): + self._log.append(("group.adopt_resident_reload", self.name)) + self.state = "mapped" + + def try_release(self): + self._log.append(("group.try_release", self.name)) + return False + + def prepare_remap(self, control_stream): + self._log.append(("group.prepare_remap", self.name)) + return True + + def adopt_reload_submission(self): + self._log.append(("group.adopt_reload_submission", self.name)) + + +class _RecordingRunner: + def __init__(self, groups, slot_bank=None): + self.local_graph_offload_groups = groups + self.local_graph_reload_state = None + self.local_graph_slot_bank = slot_bank + self.local_graph_reload_event = None + + +def _configure_ping_pong_manager(mgr, log): + mgr._local_graph_bank_d2h_events = [ + _RecordingEvent(log, "bank0_d2h"), + _RecordingEvent(log, "bank1_d2h"), + ] + mgr._local_graph_bank_d2h_recorded = [False, False] + mgr._local_graph_bank_consumed_events = [None, None] + + +def _make_local_graph_batch_manager(log): + """Build a PipelineOffloadManager with recording streams/events, bypassing __init__.""" + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + mgr = PipelineOffloadManager.__new__(PipelineOffloadManager) + mgr._d2h_stream = _RecordingStream("d2h", log) + mgr._h2d_stream = _RecordingStream("h2d", log) + mgr._local_graph_d2h_bytes = 0 + mgr._local_graph_h2d_bytes = 0 + return mgr + + +def _record_remap_batch(monkeypatch, log): + """Replace native reload submission and stream wait with recording stubs.""" + import megatron.core.pipeline_parallel.fine_grained_activation_offload as offload + + def remap_batch(allocations, host_tensors, stream): + log.append( + ( + "runner.remap_and_copy_after", + tuple(allocation[0] for allocation in allocations), + tuple(host_tensor[0] for host_tensor in host_tensors), + stream.name, + ) + ) + return object() + + def wait_on_stream(context, stream): + log.append(("runner.wait_remap_copy_on_stream", stream.name)) + + monkeypatch.setattr(offload, "remap_and_copy_after", remap_batch) + monkeypatch.setattr(offload, "wait_remap_copy_on_stream", wait_on_stream) + + +def test_local_graph_forward_replay_enqueues_without_blocking(monkeypatch): + """Forward replay only enqueues D2H on the shared stream; no wait, no release.""" + log = [] + monkeypatch.setattr(torch.cuda, "current_stream", lambda: "compute") + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + + groups = [ + _RecordingGraphGroup("expert_fc1", log), + _RecordingGraphGroup("moe_act", log), + _RecordingGraphGroup("empty_group", log, num_tensors=0), + ] + mgr = _make_local_graph_batch_manager(log) + + mgr.local_graph_forward_replay(_RecordingRunner(groups)) + + events = [entry[0] for entry in log] + # Every non-empty group enqueues in order; the empty group participates in nothing. + enqueue_names = [entry[1] for entry in log if entry[0] == "group.enqueue_d2h"] + assert enqueue_names == ["expert_fc1", "moe_act"] + # No event synchronize at forward time: the sync belongs to backward. + assert not any(e.endswith(".synchronize") for e in events) + assert "group.wait_and_release" not in events + # Byte accounting covers all groups (the empty one contributes zero). + assert mgr._local_graph_d2h_bytes == sum(g.logical_bytes for g in groups) + + +def test_local_graph_backward_prepare_finish_wait_pipeline(monkeypatch): + """Reload phases preserve reverse order and consume via one stream event wait.""" + log = [] + _record_remap_batch(monkeypatch, log) + compute_stream = _RecordingStream("compute", log) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + monkeypatch.setattr( + torch.cuda, "Event", lambda *args, **kwargs: _RecordingEvent(log, "reload_ready") + ) + + names = ["expert_fc1", "core_attn", "moe_act", "attn_proj"] + groups = [_RecordingGraphGroup(name, log) for name in names] + mgr = _make_local_graph_batch_manager(log) + runner = _RecordingRunner(groups) + + assert mgr.local_graph_backward_prepare(runner) + assert runner.local_graph_reload_state == "reload_pending" + assert mgr.local_graph_backward_finish_prepare(runner) + assert runner.local_graph_reload_state == "reload_pending" + assert mgr.local_graph_backward_wait_ready(runner, compute_stream) + + reverse_names = list(reversed(names)) + assert [e for e in log if e[0] == "runner.remap_and_copy_after"] == [ + ( + "runner.remap_and_copy_after", + tuple(reverse_names), + tuple(reverse_names), + "h2d", + ) + ] + assert [e[1] for e in log if e[0] == "group.prepare_remap"] == reverse_names + assert [e[1] for e in log if e[0] == "group.adopt_reload_submission"] == reverse_names + assert [e for e in log if e[0] == "runner.wait_remap_copy_on_stream"] == [ + ("runner.wait_remap_copy_on_stream", "compute") + ] + assert runner.local_graph_reload_state is None + assert mgr._local_graph_h2d_bytes == sum(g.logical_bytes for g in groups) + + +def test_local_graph_backward_wait_ready_primes_unprefetched_runner(monkeypatch): + """The first backward runner synchronously submits and finishes its own reload.""" + log = [] + _record_remap_batch(monkeypatch, log) + compute_stream = _RecordingStream("compute", log) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + monkeypatch.setattr( + torch.cuda, "Event", lambda *args, **kwargs: _RecordingEvent(log, "reload_ready") + ) + runner = _RecordingRunner([_RecordingGraphGroup("first", log)]) + mgr = _make_local_graph_batch_manager(log) + + assert mgr.local_graph_backward_wait_ready(runner, compute_stream) + + assert [entry[0] for entry in log] == [ + "runner.remap_and_copy_after", + "group.prepare_remap", + "runner.wait_remap_copy_on_stream", + "group.adopt_reload_submission", + ] + + +def test_local_graph_ping_pong_forward_fences_bank_reuse(monkeypatch): + """A reused resident bank waits for its previous D2H before graph replay.""" + log = [] + monkeypatch.setattr(torch.cuda, "current_stream", lambda: _RecordingStream("compute", log)) + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + mgr = _make_local_graph_batch_manager(log) + _configure_ping_pong_manager(mgr, log) + groups = [_RecordingGraphGroup("expert_fc1", log)] + first = _RecordingRunner(groups, slot_bank=0) + second = _RecordingRunner([_RecordingGraphGroup("expert_fc1", log)], slot_bank=0) + + # Simulate the first replay completing its D2H, then reuse the same bank. + mgr.local_graph_forward_replay(first) + assert mgr.local_graph_forward_wait_ready(second, _RecordingStream("replay", log)) + assert ("replay.wait_event", mgr._local_graph_bank_d2h_events[0]) in log + assert not any(entry[0] == "group.try_release" for entry in log) + + +def test_local_graph_ping_pong_backward_uses_two_banks_and_consumed_fence(monkeypatch): + """Adjacent reloads use different banks; a later same-bank reload fences consumption.""" + log = [] + compute_stream = _RecordingStream("compute", log) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + monkeypatch.setattr(torch.cuda, "Event", lambda *args, **kwargs: _RecordingEvent(log, "reload_ready")) + monkeypatch.setattr("megatron.core.pipeline_parallel.fine_grained_activation_offload.remap_and_copy_after", lambda *args: pytest.fail("VMM remap used in resident mode")) + mgr = _make_local_graph_batch_manager(log) + _configure_ping_pong_manager(mgr, log) + mgr._local_graph_bank_d2h_recorded = [True, True] + runners = [ + _RecordingRunner([_RecordingGraphGroup("g0", log)], 0), + _RecordingRunner([_RecordingGraphGroup("g1", log)], 1), + _RecordingRunner([_RecordingGraphGroup("g2", log)], 0), + ] + consumed = _RecordingEvent(log, "consumed0") + + assert mgr.local_graph_backward_prepare(runners[0]) + assert mgr.local_graph_backward_wait_ready(runners[0], compute_stream) + mgr.local_graph_backward_mark_consumed(runners[0], consumed) + assert mgr.local_graph_backward_prepare(runners[1]) + assert mgr.local_graph_backward_wait_ready(runners[1], compute_stream) + assert mgr.local_graph_backward_prepare(runners[2]) + + h2d_waits = [entry for entry in log if entry[0] == "h2d.wait_event"] + # Runner 2 reuses bank 0 and must wait for runner 0's consumed event. + assert any(entry[1] is consumed for entry in h2d_waits) + assert [entry[1] for entry in log if entry[0] == "group.enqueue_resident_h2d"] == ["g0", "g1", "g2"] + assert not any(entry[0] == "runner.remap_and_copy_after" for entry in log) + + + """With no active groups, neither replay path touches streams or groups.""" + log = [] + monkeypatch.setattr(torch.cuda, "current_stream", lambda: "compute") + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + + groups = [_RecordingGraphGroup("empty", log, num_tensors=0)] + mgr = _make_local_graph_batch_manager(log) + runner = _RecordingRunner(groups) + + mgr.local_graph_forward_replay(runner) + mgr.local_graph_backward_replay(runner) + + assert log == [] + assert mgr._local_graph_d2h_bytes == 0 + assert mgr._local_graph_h2d_bytes == 0 diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index f387e58953e..5856eb4f107 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -65,6 +65,59 @@ fp8_available, _ = check_fp8_support() +def test_backward_runner_links_follow_recorded_order(): + """Backward prefetch links use recorded schedule order and include empty runners.""" + + class Runner: + next_bwd_runner = None + + first, empty, last = Runner(), Runner(), Runner() + old_record = _CudagraphGlobalRecord.cudagraph_record + old_runners = _CudagraphGlobalRecord.backward_runners + try: + _CudagraphGlobalRecord.cudagraph_record = [ + (first, "fwd", (), {}, None), + (first, "bwd"), + (empty, "bwd"), + (last, "bwd"), + ] + _CudagraphGlobalRecord._link_backward_runners() + + assert first.next_bwd_runner is empty + assert empty.next_bwd_runner is last + assert last.next_bwd_runner is None + finally: + _CudagraphGlobalRecord.cudagraph_record = old_record + _CudagraphGlobalRecord.backward_runners = old_runners + + +def test_backward_runner_links_assign_ping_pong_banks(): + """Resident local-graph slots follow the recorded backward consumption order.""" + + class Runner: + next_bwd_runner = None + + def __init__(self, uses_local_offload): + self.uses_local_offload = uses_local_offload + + def _uses_local_graph_activation_offload(self): + return self.uses_local_offload + + runners = [Runner(True), Runner(False), Runner(True)] + old_record = _CudagraphGlobalRecord.cudagraph_record + old_runners = _CudagraphGlobalRecord.backward_runners + try: + _CudagraphGlobalRecord.cudagraph_record = [(runner, "bwd") for runner in runners] + _CudagraphGlobalRecord._link_backward_runners() + assert [runner.local_graph_slot_bank for runner in runners] == [0, None, 1] + assert runners[0].next_bwd_runner is runners[1] + assert runners[1].next_bwd_runner is runners[2] + assert runners[2].next_bwd_runner is None + finally: + _CudagraphGlobalRecord.cudagraph_record = old_record + _CudagraphGlobalRecord.backward_runners = old_runners + + def _base_cuda_graph_config(**kwargs) -> TransformerConfig: return TransformerConfig(num_layers=2, hidden_size=64, num_attention_heads=4, **kwargs) From 92f1a97259d44b3a7e1d441753b94eab8af689a3 Mon Sep 17 00:00:00 2001 From: dydy Date: Wed, 2 Sep 2026 14:05:06 +0800 Subject: [PATCH 4/6] feat: ping-pong & resdient --- .../fine_grained_activation_offload.py | 131 ++++++++++++++++-- megatron/core/transformer/cuda_graphs.py | 83 ++++++++--- .../core/transformer/transformer_config.py | 37 +++-- ...test_fine_grained_activation_offloading.py | 66 ++++++++- .../transformer/test_cuda_graphs.py | 42 ++++-- 5 files changed, 299 insertions(+), 60 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index aa5d1cecba0..eebb570a107 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -428,6 +428,11 @@ def __init__(self, name: str, min_tensor_size: int): self.remap_context = None self.slot_bank = None self.resident = False + # Build-time ownership of the resident bank. The backward graph must + # restore the values captured by this exact forward runner, not merely + # whichever runner most recently used the same ping-pong bank. + self.capture_forward_runner = None + self.capture_snapshot_event = None def _eligible(self, tensor): return ( @@ -532,6 +537,34 @@ def logical_bytes(self): def physical_bytes(self): return sum(allocation.aligned_bytes for allocation in self.allocations) + def _debug_slot_state(self, stage): + """Log VMM and tensor storage state at resident-slot checkpoints.""" + if os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_DEBUG") != "1": + return + slots = [] + for index, (allocation, tensor, host_tensor) in enumerate( + zip(self.allocations, self.device_tensors, self.host_tensors) + ): + try: + device_storage_bytes = tensor.untyped_storage().nbytes() + except Exception: # pragma: no cover - backend diagnostic only + device_storage_bytes = "unavailable" + try: + allocation_info = allocation.info() + except Exception as exc: # pragma: no cover - backend diagnostic only + allocation_info = f"error={type(exc).__name__}: {exc}" + slots.append( + f"slot={index} device_shape={tuple(tensor.shape)} " + f"device_numel={tensor.numel()} device_ptr={hex(tensor.data_ptr())} " + f"device_storage_bytes={device_storage_bytes} " + f"host_shape={tuple(host_tensor.shape)} host_numel={host_tensor.numel()} " + f"allocation={allocation_info}" + ) + PipelineOffloadManager._local_graph_debug( + f"checkpoint={stage} group={self.name} bank={self.slot_bank} state={self.state} " + + " | ".join(slots) + ) + def rebind_resident(self, manager, slot_bank, group_index): """Replace provisional capture allocations with a shared resident bank.""" if self.resident: @@ -549,22 +582,42 @@ def rebind_resident(self, manager, slot_bank, group_index): ) self.allocations.append(allocation) self.device_tensors.append(allocation.tensor) + self._debug_slot_state("after_resident_rebind") for allocation in old_allocations: allocation.close() - def capture_snapshot(self): - """Preserve capture-time values before another runner reuses this bank.""" + def capture_snapshot(self, runner=None): + """Preserve values and ownership for this runner's later bwd capture.""" if not self.resident: return + self._debug_slot_state("before_capture_snapshot") + # This copy is intentionally synchronous: it establishes the host + # snapshot before another forward is allowed to reuse the bank. for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): host_tensor.copy_(device_tensor, non_blocking=False) - - def capture_restore(self): - """Restore capture-time values before capturing the backward graph.""" + self.capture_forward_runner = runner + # Keep an explicit stream dependency even on backends where the + # synchronous copy currently happens to drain the stream. MUSA graph + # capture and the caller's stream are not assumed to be identical. + self.capture_snapshot_event = torch.cuda.Event() + self.capture_snapshot_event.record(torch.cuda.current_stream()) + self._debug_slot_state("after_capture_snapshot") + + def capture_restore(self, forward_runner=None): + """Restore the exact forward snapshot before capturing its bwd graph.""" if not self.resident: return + if self.capture_forward_runner is not forward_runner: + raise RuntimeError( + f"{self.name}: backward capture is bound to the wrong forward runner " + f"(expected={id(self.capture_forward_runner)}, actual={id(forward_runner)})" + ) + if self.capture_snapshot_event is not None: + torch.cuda.current_stream().wait_event(self.capture_snapshot_event) + self._debug_slot_state("before_capture_restore") for device_tensor, host_tensor in zip(self.device_tensors, self.host_tensors): device_tensor.copy_(host_tensor, non_blocking=False) + self._debug_slot_state("after_capture_restore") def enqueue_resident_d2h(self, d2h_stream, compute_stream): """Save a resident slot without releasing its fixed mapping.""" @@ -844,18 +897,27 @@ def local_graph_resident_allocation( return allocations[key] def local_graph_capture_snapshot(self, runner): - """Preserve shared-slot values needed by this runner's later bwd capture.""" + """Preserve the values belonging to this forward runner.""" if not self._local_graph_ping_pong_enabled(runner): return for group in runner.local_graph_offload_groups: - group.capture_snapshot() - - def local_graph_capture_restore(self, runner): - """Restore shared-slot values before this runner's backward capture.""" + group.capture_snapshot(runner) + # The next runner may reuse this bank immediately during its warmup. + # Drain every device stream here, not just the current stream, so a + # side-stream producer cannot still touch the old bank while the next + # graph is being prepared. + torch.cuda.synchronize() + + def local_graph_capture_restore(self, runner, forward_runner=None): + """Restore the forward runner explicitly bound to this bwd runner.""" if not self._local_graph_ping_pong_enabled(runner): return + if forward_runner is None: + forward_runner = runner + if forward_runner is not runner: + raise RuntimeError("local graph bwd runner must use its own forward snapshot") for group in runner.local_graph_offload_groups: - group.capture_restore() + group.capture_restore(forward_runner) def begin_local_graph_discovery(self, runner): """Discover saved activation descriptors during the final eager warmup.""" @@ -1002,6 +1064,27 @@ def local_graph_forward_replay(self, runner): for group in active_groups: group.enqueue_d2h(self.d2h_stream, compute_stream) self.drain_pending_d2h(runner) + # The first backward-consumed runner has no earlier backward replay + # to look ahead from: without this its own backward() would submit + # the remap and immediately host-wait for it, absorbing the whole + # RemapWorker latency (muMemSetAccess drains in-flight DMA, ~40 µs + # per offloaded MiB) on the critical path. Dispatch its reload now, + # right after the forward D2H submission: the RemapWorker internally + # waits on each slot's release dependency, so this stays behind the + # D2H completion while overlapping the remaining forward replay. + # Later runners keep the existing look-ahead from the predecessor's + # backward(); only the per-runner (non-bank) slot path does this. + if ( + getattr(runner, "is_first_bwd_runner", False) + and not self._local_graph_ping_pong_enabled(runner) + and getattr(runner, "local_graph_reload_state", None) is None + ): + # Async here on purpose: the chain head's release dependency is + # its own just-submitted forward D2H burst, so blocking now + # would serialize the offload copies. Its context completes + # during the remaining forward work; backward's synchronous + # scheduling takes over from the second runner on. + self.local_graph_backward_prepare(runner, block_until_submitted=False) self._local_graph_debug(f"d2h bytes={step_bytes} total_d2h={self._local_graph_d2h_bytes}") def drain_pending_d2h(self, runner): @@ -1016,8 +1099,17 @@ def _local_graph_active_groups(runner): groups = getattr(runner, "local_graph_offload_groups", []) return [group for group in reversed(groups) if group.device_tensors] - def local_graph_backward_prepare(self, runner): - """Submit one runner's resident reload or VMM remap/H2D batch.""" + def local_graph_backward_prepare(self, runner, block_until_submitted=True): + """Submit one runner's resident reload or VMM remap/H2D batch. + + With ``block_until_submitted=True`` (per-runner slot path default) the + call additionally blocks until the RemapWorker finished the VMM + transition and submitted the H2D — see the call site in + ``Graphed.backward`` for why this runs before the predecessor's + replay. The forward-time priming of the backward-chain head passes + ``False`` because its release dependency (the forward D2H burst) is + still in flight there; blocking would serialize the offload copies. + """ active_groups = self._local_graph_active_groups(runner) if not active_groups: return False @@ -1055,6 +1147,19 @@ def local_graph_backward_prepare(self, runner): reload_context = remap_and_copy_after( allocations, host_tensors, self.h2d_stream ) + if block_until_submitted: + # Synchronous remap: block right here until the worker finished + # the VMM transition and submitted the H2D, instead of deferring + # the host wait to this runner's own backward() boundary. Call + # this BEFORE the predecessor's backward replay is submitted: + # the GPU is momentarily quiet, so muMemSetAccess drains almost + # no in-flight DMA, and the reload H2D is already queued on the + # h2d stream by the time the backward graph launches — compute + # and H2D then run concurrently instead of the H2D landing in + # the bubble after the compute segment. The consumer-side wait + # in wait_ready afterwards only installs the compute-stream + # dependency on an already-completed context. + wait_remap_copy_on_stream(reload_context, self.h2d_stream) for group in active_groups: group.prepare_remap(reload_context) runner.local_graph_reload_context = reload_context diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 6877ac81bab..feb9540d972 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -30,7 +30,10 @@ ) from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule -from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_config import ( + LOCAL_GRAPH_ACTIVATION_OFFLOAD_MODULES, + TransformerConfig, +) from megatron.core.utils import ( get_attr_wrapped_model, get_torch_version, @@ -528,6 +531,10 @@ def record_fwd_graph(cls, runner, args, kwargs, out): @classmethod def record_bwd_graph(cls, runner): """Record a bwd graph to 'cudagraph_record""" + # The autograd node is created by this runner's forward invocation. + # Preserve that identity explicitly; resident banks may be reused by a + # later forward runner before this bwd graph is captured. + runner.local_graph_forward_runner = runner cls.cudagraph_record.append((runner, "bwd")) @classmethod @@ -541,12 +548,27 @@ def _link_backward_runners(cls): runner, "_uses_local_graph_activation_offload", lambda: False )() runner.local_graph_slot_bank = ( - local_graph_offload_index % 2 if uses_local_offload else None + local_graph_offload_index % 2 + if uses_local_offload and os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_PING_PONG", "1") != "0" + else None ) if uses_local_offload: local_graph_offload_index += 1 for runner, next_runner in zip(cls.backward_runners, cls.backward_runners[1:]): runner.next_bwd_runner = next_runner + # The head of the backward chain is consumed first, yet no earlier + # backward replay exists to look ahead from. Mark it so its reload can + # be dispatched right after its forward D2H instead of inside its own + # backward() host wait (which would absorb the whole RemapWorker latency + # including the muMemSetAccess DMA-drain). + for runner in cls.backward_runners: + runner.is_first_bwd_runner = False + # A bwd graph is captured from this runner's forward graph state. + # Keep that association explicit even when adjacent runners share + # a resident ping-pong bank. + runner.local_graph_forward_runner = runner + if cls.backward_runners: + cls.backward_runners[0].is_first_bwd_runner = True @classmethod def create_cudagraphs(cls): @@ -731,6 +753,8 @@ def delete_cuda_graphs(): for runner in _CudagraphGlobalRecord.backward_runners: runner.next_bwd_runner = None runner.local_graph_reload_state = None + runner.is_first_bwd_runner = False + runner.local_graph_forward_runner = runner # Reset global tracking state _CudagraphGlobalRecord.cudagraph_created = False @@ -878,30 +902,44 @@ def backward(ctx, *grads): offload_manager = None next_runner = runner.next_bwd_runner - next_prepared = False if runner.local_graph_offload_groups: offload_manager = runner._local_graph_offload_manager() replay_stream = runner.stream if runner.use_stream else torch.cuda.current_stream() offload_manager.local_graph_backward_wait_ready(runner, replay_stream) + # Submit the successor's reload BEFORE this runner's backward replay: + # with synchronous remap, prepare() blocks until VMM transition and H2D + # submission are done. Doing it pre-replay means muMemSetAccess runs + # while the GPU is quiet (the predecessor graph just finished), and the + # successor's H2D is already on the h2d stream when this backward graph + # launches — current-layer compute and next-layer reload H2D overlap on + # the device. Resident ping-pong keeps its async one-layer semantics: + # prepare is cheap there, and deferring past replay preserves overlap. if next_runner is not None and next_runner.local_graph_offload_groups: if offload_manager is None: offload_manager = next_runner._local_graph_offload_manager() - next_prepared = offload_manager.local_graph_backward_prepare(next_runner) - - try: - if runner.use_stream: - runner.stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(runner.stream): - runner.bwd_graph.replay() - torch.cuda.current_stream().wait_event(runner.bwd_completion_event) + if offload_manager._local_graph_ping_pong_enabled(next_runner): + offload_manager.local_graph_backward_prepare(next_runner) else: + second_runner = next_runner.next_bwd_runner + offload_manager.local_graph_backward_prepare(next_runner) + # Prime the layer after next as well: its synchronous remap now + # runs while only the (already submitted) next-layer H2D is in + # flight, keeping the remap chain ahead without ever parking a + # context behind a running communication/compute segment. + if ( + second_runner is not None + and second_runner.local_graph_offload_groups + and not offload_manager._local_graph_ping_pong_enabled(second_runner) + ): + offload_manager.local_graph_backward_prepare(second_runner) + + if runner.use_stream: + runner.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(runner.stream): runner.bwd_graph.replay() - finally: - # The current graph has already been submitted. Waiting for the - # successor's remap worker and enqueueing H2D now overlaps both with - # current-runner GPU execution. - if next_prepared: - offload_manager.local_graph_backward_finish_prepare(next_runner) + torch.cuda.current_stream().wait_event(runner.bwd_completion_event) + else: + runner.bwd_graph.replay() runner.bwd_graph_replay_complete_event.record(torch.cuda.current_stream()) if runner.local_graph_offload_groups: @@ -971,6 +1009,8 @@ def __init__( self.local_graph_slot_bank = None self.local_graph_reload_state = None self.local_graph_reload_event = None + self.is_first_bwd_runner = False + self.local_graph_forward_runner = self self.fwd_graph_recorded = False self.bwd_graph_recorded = False @@ -1109,15 +1149,14 @@ def __str__(self): ) def _uses_local_graph_activation_offload(self): - """Whether this runner captures the whole layer with MoE-only activation offload.""" + """Whether this runner captures a whole layer with local Graph activation offload.""" config = getattr(self.base_module, "config", None) return bool( config is not None and config.cuda_graph_impl == "local" and not config.cuda_graph_modules and config.fine_grained_activation_offloading - and set(config.offload_modules or []) - <= {"expert_fc1", "moe_act", "fused_group_mlp"} + and set(config.offload_modules or []) <= LOCAL_GRAPH_ACTIVATION_OFFLOAD_MODULES ) @staticmethod @@ -1549,7 +1588,9 @@ def create_bwd_graph(self): # Shared ping-pong slots may have been overwritten while later forward # graphs were captured. Restore this runner before capturing its consumers. if self.local_graph_offload_groups: - self._local_graph_offload_manager().local_graph_capture_restore(self) + self._local_graph_offload_manager().local_graph_capture_restore( + self, self.local_graph_forward_runner + ) # Freeze GC, to speed up capture time ~15-20x. if FREEZE_GC: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1dd19ec043d..ca45256c2da 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -27,6 +27,12 @@ ) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout +# Activation groups whose saved tensors can use fixed addresses during whole-layer +# local CUDA graph capture. Keep this list shared with the runner predicate. +LOCAL_GRAPH_ACTIVATION_OFFLOAD_MODULES = frozenset( + {"expert_fc1", "moe_act", "fused_group_mlp", "attn_norm", "qkv_linear"} +) + from .._rank_utils import log_single_rank from ..fusions.fused_bias_geglu import quick_gelu from ..model_parallel_config import ModelParallelConfig @@ -1895,15 +1901,15 @@ def __post_init__(self): self.delta_offload_bytes_across_pp_ranks >= 0 ), "delta_offload_bytes_across_pp_ranks must be non-negative." - if "fused_group_mlp" in self.offload_modules: - if not self.use_transformer_engine_op_fuser: - raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.") - moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules) - if moe_partial_offload: - raise ValueError( - "fused_group_mlp offloads the whole fused grouped MLP and cannot be " - f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}" - ) + # if "fused_group_mlp" in self.offload_modules: + # if not self.use_transformer_engine_op_fuser: + # raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.") + # moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules) + # if moe_partial_offload: + # raise ValueError( + # "fused_group_mlp offloads the whole fused grouped MLP and cannot be " + # f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}" + # ) if self.moe_paged_stash: if self.cpu_offloading: raise ValueError("moe_paged_stash cannot be enabled with cpu_offloading.") @@ -2599,12 +2605,15 @@ def _scope_to_str(s): if self.fine_grained_activation_offloading: offload_modules = set(self.offload_modules or []) if self.cuda_graph_impl == "local": - local_supported_offload_modules = {"expert_fc1", "moe_act", "fused_group_mlp"} - unsupported_offload_modules = offload_modules - local_supported_offload_modules + supported_offload_modules = ( + LOCAL_GRAPH_ACTIVATION_OFFLOAD_MODULES + if not self.cuda_graph_modules + else {"expert_fc1", "moe_act", "fused_group_mlp"} + ) + unsupported_offload_modules = offload_modules - supported_offload_modules assert not unsupported_offload_modules, ( "fine-grained activation offloading with cuda_graph_impl='local' " - "only supports offload_modules 'expert_fc1', 'moe_act', and " - "'fused_group_mlp'. " + "only supports the selected activation modules for this graph scope. " f"Unsupported offload_modules: {sorted(unsupported_offload_modules)}." ) local_moe_offload = ( @@ -2613,7 +2622,7 @@ def _scope_to_str(s): and offload_modules <= {"expert_fc1", "moe_act", "fused_group_mlp"} ) assert ( - self.cuda_graph_impl in ("transformer_engine", "full_iteration") + self.cuda_graph_impl in ("transformer_engine", "full_iteration", "local") or local_moe_offload ), ( "fine-grained activation offloading is only supported with " diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index bdf0aa1b3b3..dbadd98efbf 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -1159,14 +1159,23 @@ def test_local_graph_backward_prepare_finish_wait_pipeline(monkeypatch): assert [e[1] for e in log if e[0] == "group.prepare_remap"] == reverse_names assert [e[1] for e in log if e[0] == "group.adopt_reload_submission"] == reverse_names assert [e for e in log if e[0] == "runner.wait_remap_copy_on_stream"] == [ - ("runner.wait_remap_copy_on_stream", "compute") + # Synchronous remap: prepare blocks on the h2d stream first, then the + # consumer wait installs the compute-stream dependency on the + # already-completed context. + ("runner.wait_remap_copy_on_stream", "h2d"), + ("runner.wait_remap_copy_on_stream", "compute"), ] assert runner.local_graph_reload_state is None assert mgr._local_graph_h2d_bytes == sum(g.logical_bytes for g in groups) def test_local_graph_backward_wait_ready_primes_unprefetched_runner(monkeypatch): - """The first backward runner synchronously submits and finishes its own reload.""" + """The unprefetched fallback path synchronously submits and waits its own reload. + + prepare() blocks until the worker submitted the H2D (wait before + prepare_remap bookkeeping); wait_ready then only installs the + compute-stream dependency on the completed context. + """ log = [] _record_remap_batch(monkeypatch, log) compute_stream = _RecordingStream("compute", log) @@ -1181,8 +1190,61 @@ def test_local_graph_backward_wait_ready_primes_unprefetched_runner(monkeypatch) assert mgr.local_graph_backward_wait_ready(runner, compute_stream) assert [entry[0] for entry in log] == [ + "runner.remap_and_copy_after", + "runner.wait_remap_copy_on_stream", + "group.prepare_remap", + "runner.wait_remap_copy_on_stream", + "group.adopt_reload_submission", + ] + + +def test_local_graph_forward_replay_primes_first_backward_runner(monkeypatch): + """The backward-chain head dispatches its reload at forward-D2H time. + + Without this, the runner consumed first in backward has no predecessor + replay to look ahead from, so its own backward() would submit the remap + and immediately host-wait for it (absorbing the RemapWorker latency, + including the muMemSetAccess in-flight-DMA drain) on the critical path. + """ + log = [] + _record_remap_batch(monkeypatch, log) + compute_stream = _RecordingStream("compute", log) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda s: nullcontext()) + monkeypatch.setattr( + torch.cuda, "Event", lambda *args, **kwargs: _RecordingEvent(log, "reload_ready") + ) + first = _RecordingRunner([_RecordingGraphGroup("first", log)]) + first.is_first_bwd_runner = True + later = _RecordingRunner([_RecordingGraphGroup("later", log)]) + mgr = _make_local_graph_batch_manager(log) + + # Forward replay of the chain head submits its reload immediately after + # the D2H burst; later runners are not touched here. + mgr.local_graph_forward_replay(first) + assert first.local_graph_reload_state == "reload_pending" + assert [e[0] for e in log if e[0] == "runner.remap_and_copy_after"] == [ + "runner.remap_and_copy_after" + ] + assert not any(e[1] == "later" for e in log if e[0] == "group.prepare_remap") + + mgr.local_graph_forward_replay(later) + assert later.local_graph_reload_state is None + assert [e[0] for e in log if e[0] == "runner.remap_and_copy_after"] == [ + "runner.remap_and_copy_after" + ] + + # Backward consumption of the head no longer submits anything; it only + # installs the stream wait and adopts the already-submitted reload. + # (`later`'s D2H entries sit between because its forward replay ran too.) + assert mgr.local_graph_backward_wait_ready(first, compute_stream) + assert [entry[0] for entry in log] == [ + "group.enqueue_d2h", + "group.try_release", "runner.remap_and_copy_after", "group.prepare_remap", + "group.enqueue_d2h", + "group.try_release", "runner.wait_remap_copy_on_stream", "group.adopt_reload_submission", ] diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 5856eb4f107..d495ebd2f30 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -35,6 +35,7 @@ CudaGraphManager, TECudaGraphHelper, _CudagraphGlobalRecord, + _CudaGraphRunner, create_cudagraphs, ) from megatron.core.transformer.enums import ( @@ -113,6 +114,12 @@ def _uses_local_graph_activation_offload(self): assert runners[0].next_bwd_runner is runners[1] assert runners[1].next_bwd_runner is runners[2] assert runners[2].next_bwd_runner is None + # Only the first backward-consumed runner is marked; its reload is + # dispatched at forward-D2H time instead of inside its own backward(). + assert runners[0].is_first_bwd_runner is True + assert runners[1].is_first_bwd_runner is False + assert runners[2].is_first_bwd_runner is False + assert all(runner.local_graph_forward_runner is runner for runner in runners) finally: _CudagraphGlobalRecord.cudagraph_record = old_record _CudagraphGlobalRecord.backward_runners = old_runners @@ -122,6 +129,22 @@ def _base_cuda_graph_config(**kwargs) -> TransformerConfig: return TransformerConfig(num_layers=2, hidden_size=64, num_attention_heads=4, **kwargs) +def test_runner_recognizes_attention_norm_local_graph_offload(): + """Whole-layer local graphs treat attention norm as a resident activation group.""" + + class Module: + config = _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[], + fine_grained_activation_offloading=True, + offload_modules=['attn_norm'], + ) + + runner = object.__new__(_CudaGraphRunner) + runner.base_module = Module() + assert runner._uses_local_graph_activation_offload() + + def _validated_cuda_graph_cli_args(monkeypatch, cli_args=None, **overrides): destroy_global_vars() destroy_num_microbatches_calculator() @@ -187,16 +210,15 @@ def test_local_impl_rejects_unsupported_activation_offload_scope(self): offload_modules=['qkv_linear'], ) - def test_local_impl_rejects_full_layer_graph_with_activation_offload(self): - with pytest.raises( - AssertionError, match="not supported with whole-layer CUDA graph capture" - ): - _base_cuda_graph_config( - cuda_graph_impl='local', - cuda_graph_modules=[], - fine_grained_activation_offloading=True, - offload_modules=['expert_fc1'], - ) + def test_local_impl_allows_full_layer_graph_with_activation_offload(self): + cfg = _base_cuda_graph_config( + cuda_graph_impl='local', + cuda_graph_modules=[], + fine_grained_activation_offloading=True, + offload_modules=['attn_norm'], + ) + assert cfg.cuda_graph_impl == 'local' + assert cfg.offload_modules == ['attn_norm'] def test_local_impl_rejects_moe_router_graph_with_mlp_norm_offload(self): with pytest.raises( From c799b76e71db9b1750313907e10e12be429f553d Mon Sep 17 00:00:00 2001 From: dydy Date: Thu, 17 Sep 2026 09:30:57 +0800 Subject: [PATCH 5/6] todo: resdient slot --- .../fine_grained_activation_offload.py | 584 ++++++++++++++---- megatron/core/transformer/cuda_graphs.py | 58 +- ...test_fine_grained_activation_offloading.py | 13 +- 3 files changed, 488 insertions(+), 167 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index eebb570a107..f77025c94ae 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os +import time from collections import defaultdict, deque from contextlib import nullcontext from typing import Any, Dict, Optional, Tuple @@ -17,14 +18,20 @@ try: from transformer_engine.pytorch.vmm_activation import ( + launch_remap_slot_h2d, release_hooks_after, - remap_and_copy_after, - wait_remap_copy_on_stream, + remap_and_copy_slot_after, + remap_only_slot_after, + wait_remap_slot_on_stream, + wait_until_remap_slot_submitted, ) except ImportError: release_hooks_after = None # eager-only fallback; local graph paths require TE - remap_and_copy_after = None - wait_remap_copy_on_stream = None + remap_and_copy_slot_after = None + remap_only_slot_after = None + launch_remap_slot_h2d = None + wait_remap_slot_on_stream = None + wait_until_remap_slot_submitted = None def debug_rank(message): @@ -420,11 +427,17 @@ def __init__(self, name: str, min_tensor_size: int): self.device_tensors = [] self.host_tensors = [] self.capture_index = 0 + self.capture_slots = [] + self.consumer_slots = [] + self.capture_epoch = 0 + self.capture_group_index = None self.state = "discovering" + self.d2h_event = torch.cuda.Event() self.h2d_event = torch.cuda.Event() - # Async release/remap contexts returned by the batched host functions. + # Async release/remap contexts returned by the per-slot host functions. self.release_context = None + self.release_contexts = [] self.remap_context = None self.slot_bank = None self.resident = False @@ -459,9 +472,31 @@ def discover_tensor(self, tensor): self.descriptors.append(self._descriptor(tensor)) return tensor - @staticmethod - def load_tensor(tensor): - """Return the tensor represented by a saved-tensor hook payload.""" + def load_tensor(self, tensor): + """Record the activation slot consumed by backward replay.""" + if self._eligible(tensor): + slot = next( + (item for item in self.device_tensors if item.data_ptr() == tensor.data_ptr()), + None, + ) + if slot is not None: + slot_index = self.device_tensors.index(slot) + self.consumer_slots.append(slot_index) + context = PipelineOffloadManager.get_instance()._local_graph_consumer_context + if context is not None: + context["slots"].append({ + "group_index": getattr(self, "capture_group_index", None), + "slot_index": slot_index, + "vmm_address": tensor.data_ptr(), + }) + runner = getattr(self, "capture_runner", None) + PipelineOffloadManager._local_graph_debug( + f"consume rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1} " + f"runner_id={id(runner) if runner is not None else None} " + f"group_index={getattr(self, 'capture_group_index', None)} " + f"group={self.name} slot={slot_index} " + f"vmm_va={hex(tensor.data_ptr())} epoch={self.capture_epoch}" + ) return tensor def allocate(self, manager=None, slot_bank=None, group_index=None): @@ -486,6 +521,11 @@ def allocate(self, manager=None, slot_bank=None, group_index=None): else: allocation = MUSAActivationVMMAllocation(shape, stride, dtype, device) self.allocations.append(allocation) + allocation.set_slot_id( + f"rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1}/" + f"group={group_index if group_index is not None else 'na'}/slot={tensor_index}/" + f"va=0x{allocation.address:x}" + ) self.device_tensors.append(allocation.tensor) # Host storage remains runner-local even when the GPU slot is shared. self.host_tensors.append( @@ -497,6 +537,9 @@ def begin_capture(self): """Reset ordered descriptor consumption before forward graph capture.""" assert self.state == "mapped" self.capture_index = 0 + self.capture_slots = [] + self.consumer_slots = [] + self.capture_epoch += 1 def save_tensor(self, tensor): """Capture a D2D write and save the matching VMM tensor for backward.""" @@ -513,6 +556,14 @@ def save_tensor(self, tensor): ) slot_tensor = self.device_tensors[self.capture_index] slot_tensor.copy_(tensor) + self.capture_slots.append( + { + "slot_index": self.capture_index, + "descriptor": descriptor, + "vmm_address": int(slot_tensor.data_ptr()), + "aligned_bytes": self.allocations[self.capture_index].aligned_bytes, + } + ) PipelineOffloadManager._local_graph_debug( f"capture group={self.name} slot={self.capture_index} " f"source={hex(tensor.data_ptr())} vmm_va={hex(slot_tensor.data_ptr())} " @@ -581,6 +632,11 @@ def rebind_resident(self, manager, slot_bank, group_index): slot_bank, self.name, group_index, tensor_index, shape, stride, dtype, device ) self.allocations.append(allocation) + allocation.set_slot_id( + f"rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1}/" + f"group={group_index if group_index is not None else 'na'}/slot={tensor_index}/" + f"va=0x{allocation.address:x}" + ) self.device_tensors.append(allocation.tensor) self._debug_slot_state("after_resident_rebind") for allocation in old_allocations: @@ -644,24 +700,72 @@ def adopt_resident_reload(self): raise RuntimeError(f"{self.name}: resident reload is in state {self.state}") self.state = "mapped" - def enqueue_d2h(self, d2h_stream, compute_stream): + def enqueue_d2h( + self, d2h_stream, compute_stream, release_stream=None, graph_done_event=None + ): """Queue this group's D2H copies and hand off release to a stream host func. - A single host function follows all copies on the D2H stream: once the - burst completes it moves the physical unmap/release to a resident - worker thread (driver calls are not allowed on the host-func callback - thread). No event wait or polling is needed afterwards. + The copies run on ``d2h_stream`` (ordered behind the producing compute + via ``wait_stream``) and an event records the burst's completion. The + release host func is gated ONLY on that event: it is enqueued on + ``release_stream`` after a ``wait_event(d2h_event)``, so the worker + unmaps this group as soon as ITS OWN copies finish - not when every + later group's D2H on the shared stream also completes. Driver calls + still run on the resident worker thread, not the callback thread. """ if not self.device_tensors: return if self.state != "mapped": raise RuntimeError(f"{self.name}: cannot offload slots in state {self.state}") - d2h_stream.wait_stream(compute_stream) - with torch.cuda.stream(d2h_stream): - for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): - host_tensor.copy_(device_tensor, non_blocking=True) - self.d2h_event.record(d2h_stream) - self.release_context = release_hooks_after(self.allocations, d2h_stream) + if os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_DEBUG") == "1": + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + print( + f"[vmm-offload][rank={rank}] submit_d2h group={self.name} " + f"state={self.state} allocations={len(self.allocations)} " + f"bytes={self.logical_bytes} compute_stream={getattr(compute_stream, 'musa_stream', None)} " + f"d2h_stream={getattr(d2h_stream, 'musa_stream', None)} " + f"release_stream={getattr(release_stream, 'musa_stream', None)}", + flush=True, + ) + # Event-gated per-group release: fall back to the old shared-stream + # placement only when no dedicated release stream was provided. + if release_stream is not None: + with torch.cuda.stream(d2h_stream): + if graph_done_event is not None: + d2h_stream.wait_event(graph_done_event) + else: + d2h_stream.wait_stream(compute_stream) + previous_slot_event = None + for slot_index, (host_tensor, device_tensor) in enumerate( + zip(self.host_tensors, self.device_tensors) + ): + if previous_slot_event is not None: + d2h_stream.wait_event(previous_slot_event) + host_tensor.copy_(device_tensor, non_blocking=True) + slot_done = torch.cuda.Event() + slot_done.record(d2h_stream) + previous_slot_event = slot_done + PipelineOffloadManager._local_graph_debug( + f"d2h_slot_submit ts_ns={time.time_ns()} " + f"group={self.name} slot={slot_index}" + ) + self.d2h_event.record(d2h_stream) + with torch.cuda.stream(release_stream): + release_stream.wait_event(self.d2h_event) + self.release_contexts = [] + for allocation in self.allocations: + self.release_contexts.append(allocation.release_hook_after(release_stream)) + self.release_context = self.release_contexts[-1] if self.release_contexts else None + else: + d2h_stream.wait_stream(compute_stream) + with torch.cuda.stream(d2h_stream): + for host_tensor, device_tensor in zip(self.host_tensors, self.device_tensors): + host_tensor.copy_(device_tensor, non_blocking=True) + self.d2h_event.record(d2h_stream) + self.release_contexts = [] + for allocation in self.allocations: + self.release_contexts.append(allocation.release_hook_after(d2h_stream)) + self.release_context = self.release_contexts[-1] if self.release_contexts else None self.state = "d2h_pending" def try_release(self): @@ -683,65 +787,79 @@ def wait_and_release(self): return self._adopt_released() def _release_context_done(self): - """Whether the worker finished the async release (host-side check).""" + """Whether every per-slot release request finished on the worker.""" + contexts = self.release_contexts + if contexts: + return all(dict(context.status())["done"] == 1 for context in contexts) if self.release_context is None: - raise RuntimeError(f"{self.name}: missing async release context in state {self.state}") + return all(allocation.async_release_done() for allocation in self.allocations) return dict(self.release_context.status())["done"] == 1 def _wait_release_context(self): - """Block until the worker finished; surface any worker error.""" + """Block until every per-slot release request finishes.""" + if self.release_contexts: + for context in self.release_contexts: + status = dict(context.status()) + if status["done"] != 1: + # The allocation owns the context and provides the blocking + # compatibility API without synchronizing a MUSA stream. + pass for allocation in self.allocations: allocation.wait_for_async_release() def _adopt_released(self): """Apply the worker-completed release to this group's bookkeeping.""" - status = dict(self.release_context.status()) + contexts = self.release_contexts + self.release_contexts = [] + if contexts: + for context in contexts: + status = dict(context.status()) + if status["error"]: + raise RuntimeError( + f"{self.name}: async release failed on worker: MUresult {status['error_code']}" + ) + elif self.release_context is not None: + status = dict(self.release_context.status()) + if status["error"]: + raise RuntimeError( + f"{self.name}: async release failed on worker: MUresult {status['error_code']}" + ) self.release_context = None - if status["error"]: - raise RuntimeError( - f"{self.name}: async release failed on worker: MUresult {status['error_code']}" - ) + for allocation in self.allocations: + if not allocation.async_release_done(): + allocation.wait_for_async_release() self.state = "unmapped" PipelineOffloadManager._local_graph_debug( f"release group={self.name} logical_bytes={self.logical_bytes} " f"physical_bytes={self.physical_bytes}" ) return True - def prepare_remap(self, remap_context): - """Bind this group to its runner-level asynchronous remap batch.""" + """Bind a slot request context to this group's bookkeeping.""" if not self.device_tensors: return False - if self.state not in ("d2h_pending", "unmapped"): + if self.state not in ("d2h_pending", "unmapped", "remap_pending"): raise RuntimeError(f"{self.name}: cannot prepare reload in state {self.state}") self.remap_context = remap_context self.state = "remap_pending" return True def adopt_reload_submission(self): - """Adopt remaps after the runner installed its GPU-side H2D wait.""" + """Compatibility adoption for callers submitting a complete context.""" if self.state != "remap_pending": raise RuntimeError(f"{self.name}: cannot adopt reload in state {self.state}") - for allocation in self.allocations: - allocation.adopt_async_remap() - status = dict(self.remap_context.status()) + if self.remap_context is not None: + for allocation in self.allocations: + allocation.adopt_async_remap() + status = dict(self.remap_context.status()) + if status["error"]: + raise RuntimeError( + f"{self.name}: async reload failed on worker: error {status['error_code']}" + ) self.remap_context = None self.release_context = None - if status["error"]: - raise RuntimeError( - f"{self.name}: async reload failed on worker: error {status['error_code']}" - ) self.state = "mapped" - PipelineOffloadManager._local_graph_debug( - f"reload submitted group={self.name} logical_bytes={self.logical_bytes} " - f"physical_bytes={self.physical_bytes}" - ) - - @staticmethod - def _driver_free_memory(): - from transformer_engine.pytorch.vmm_activation import vmm_driver_memory_info - - return vmm_driver_memory_info()["free_bytes"] + return True class PipelineOffloadManager: @@ -810,6 +928,11 @@ def __init__(self): self._local_graph_bank_d2h_events = [torch.cuda.Event(), torch.cuda.Event()] self._local_graph_bank_d2h_recorded = [False, False] self._local_graph_bank_consumed_events = [None, None] + # Dedicated stream for the graph-level release host func. Putting the + # host func on the shared d2h stream would serialize the next graph's + # D2H copies behind this graph's unmap callback. + self._local_graph_release_stream = torch.cuda.Stream() + self._local_graph_consumer_context = None self.do_offload = True @@ -860,8 +983,8 @@ def in_local_graph_capture(self): def _local_graph_debug(message): if os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_DEBUG") != "1": return - if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: - print(f"[local-full-graph-offload] {message}", flush=True) + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + print(f"[local-full-graph-offload][rank={rank}] {message}", flush=True) @staticmethod def _local_graph_group_bytes(groups): @@ -962,10 +1085,19 @@ def begin_local_graph_capture(self, runner): self._local_graph_capture_runner = runner self._local_graph_mode = "capture" self._local_graph_capture_groups = groups - self._local_graph_capture_group = None self._local_graph_group_index = 0 - for group in groups: - group.begin_capture() + runner.local_graph_capture_epoch = getattr(runner, "local_graph_capture_epoch", 0) + 1 + runner.local_graph_capture_metadata = { + "epoch": runner.local_graph_capture_epoch, + "rank": torch.distributed.get_rank() if torch.distributed.is_initialized() else -1, + "runner_id": id(runner), + "groups": [], + "consumer_slots": [], + } + for group_index, group in enumerate(groups): + group.capture_runner = runner + group.capture_group_index = group_index + self._local_graph_capture_group = None def end_local_graph_capture(self, success=True): """Validate capture-time descriptor consumption and retain runner slots.""" @@ -979,10 +1111,20 @@ def end_local_graph_capture(self, success=True): raise RuntimeError( f"local graph capture consumed {self._local_graph_group_index}/{len(groups)} groups" ) - for group in groups: + for group_index, group in enumerate(groups): group.finish_capture() + metadata = getattr(runner, "local_graph_capture_metadata", None) + if metadata is not None: + metadata["groups"].append( + { + "group_index": group_index, + "name": group.name, + "slots": list(group.capture_slots), + } + ) self._local_graph_debug( - f"captured groups={len(groups)} bytes={self._local_graph_group_bytes(groups)}" + f"captured groups={len(groups)} bytes={self._local_graph_group_bytes(groups)} " + f"consumer_slots={runner.local_graph_capture_metadata['consumer_slots']}" ) else: runner.local_graph_offload_groups = [] @@ -1041,8 +1183,8 @@ def local_graph_forward_wait_ready(self, runner, replay_stream): replay_stream.wait_event(consumed_event) return True - def local_graph_forward_replay(self, runner): - """Enqueue all groups' D2H after forward graph replay.""" + def local_graph_forward_replay(self, runner, graph_done_event=None): + """Enqueue D2H only after the replay completion fence.""" compute_stream = torch.cuda.current_stream() groups = getattr(runner, "local_graph_offload_groups", []) active_groups = [group for group in groups if group.device_tensors] @@ -1059,33 +1201,76 @@ def local_graph_forward_replay(self, runner): self._local_graph_bank_d2h_events[bank].record(self.d2h_stream) self._local_graph_bank_d2h_recorded[bank] = True else: - # Submit every group's D2H back-to-back on the shared stream, then - # return without a host-side wait. - for group in active_groups: - group.enqueue_d2h(self.d2h_stream, compute_stream) + # After this graph's forward replay, offload every slot in every + # group that belongs to the graph as one D2H burst and one batched + # release host func. Per-slot host funcs previously unmapped each + # allocation independently as soon as its own copy finished. + self._enqueue_graph_groups_offload( + runner, active_groups, compute_stream, graph_done_event + ) self.drain_pending_d2h(runner) - # The first backward-consumed runner has no earlier backward replay - # to look ahead from: without this its own backward() would submit - # the remap and immediately host-wait for it, absorbing the whole - # RemapWorker latency (muMemSetAccess drains in-flight DMA, ~40 µs - # per offloaded MiB) on the critical path. Dispatch its reload now, - # right after the forward D2H submission: the RemapWorker internally - # waits on each slot's release dependency, so this stays behind the - # D2H completion while overlapping the remaining forward replay. - # Later runners keep the existing look-ahead from the predecessor's - # backward(); only the per-runner (non-bank) slot path does this. - if ( - getattr(runner, "is_first_bwd_runner", False) - and not self._local_graph_ping_pong_enabled(runner) - and getattr(runner, "local_graph_reload_state", None) is None - ): - # Async here on purpose: the chain head's release dependency is - # its own just-submitted forward D2H burst, so blocking now - # would serialize the offload copies. Its context completes - # during the remaining forward work; backward's synchronous - # scheduling takes over from the second runner on. - self.local_graph_backward_prepare(runner, block_until_submitted=False) - self._local_graph_debug(f"d2h bytes={step_bytes} total_d2h={self._local_graph_d2h_bytes}") + + def _enqueue_graph_groups_offload(self, runner, active_groups, compute_stream, graph_done_event): + """D2H then batch-release every slot owned by this graph's offload groups.""" + if release_hooks_after is None: + raise RuntimeError( + "TE release_hooks_after is required to batch-offload local graph activation slots" + ) + allocations = [] + for group in active_groups: + if group.state != "mapped": + raise RuntimeError(f"{group.name}: cannot offload slots in state {group.state}") + allocations.extend(group.allocations) + if not allocations: + return + + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + self._local_graph_debug( + f"d2h_submit_begin ts_ns={time.time_ns()} rank={rank} " + f"runner_id={id(runner)} groups={[group.name for group in active_groups]} " + f"slots={len(allocations)}" + ) + if os.getenv("MEGATRON_LOCAL_GRAPH_OFFLOAD_DEBUG") == "1": + print( + f"[vmm-offload][rank={rank}] submit_graph_d2h " + f"runner_id={id(runner)} groups={[group.name for group in active_groups]} " + f"allocations={len(allocations)} " + f"compute_stream={getattr(compute_stream, 'musa_stream', None)} " + f"d2h_stream={getattr(self.d2h_stream, 'musa_stream', None)} " + f"release_stream={getattr(self._local_graph_release_stream, 'musa_stream', None)}", + flush=True, + ) + + with torch.cuda.stream(self.d2h_stream): + if graph_done_event is not None: + self.d2h_stream.wait_event(graph_done_event) + else: + self.d2h_stream.wait_stream(compute_stream) + for group in active_groups: + for slot_index, (host_tensor, device_tensor) in enumerate( + zip(group.host_tensors, group.device_tensors) + ): + host_tensor.copy_(device_tensor, non_blocking=True) + self._local_graph_debug( + f"d2h_slot_submit ts_ns={time.time_ns()} " + f"group={group.name} slot={slot_index}" + ) + group.d2h_event.record(self.d2h_stream) + + with torch.cuda.stream(self._local_graph_release_stream): + self._local_graph_release_stream.wait_event(active_groups[-1].d2h_event) + batch_context = release_hooks_after(allocations, self._local_graph_release_stream) + + for group in active_groups: + group.release_contexts = [batch_context] + group.release_context = batch_context + group.state = "d2h_pending" + + self._local_graph_debug( + f"d2h_submit_return ts_ns={time.time_ns()} rank={rank} " + f"runner_id={id(runner)} slots={len(allocations)}" + ) + def drain_pending_d2h(self, runner): """Release VMM backing for groups whose D2H already completed; never blocks.""" @@ -1140,75 +1325,197 @@ def local_graph_backward_prepare(self, runner, block_until_submitted=True): # Submit in backward consumption order. The last forward-offloaded # group has the latest release dependency, so waiting on it first # makes every earlier dependency a fast path. - allocations = [ - allocation for group in active_groups for allocation in group.allocations - ] - host_tensors = [tensor for group in active_groups for tensor in group.host_tensors] - reload_context = remap_and_copy_after( - allocations, host_tensors, self.h2d_stream + # Non-ping-pong reload is deliberately deferred until this runner is + # consumed. Creating one batch context here used to enqueue every + # slot early, allowing fast ranks to reload the whole runner long + # before graph replay. Keep the slot/group description on the + # runner; submit each group from local_graph_backward_wait_ready(). + runner.local_graph_reload_context = None + runner.local_graph_reload_groups = tuple(active_groups) + runner.local_graph_reload_slot_count = sum( + len(group.allocations) for group in active_groups ) - if block_until_submitted: - # Synchronous remap: block right here until the worker finished - # the VMM transition and submitted the H2D, instead of deferring - # the host wait to this runner's own backward() boundary. Call - # this BEFORE the predecessor's backward replay is submitted: - # the GPU is momentarily quiet, so muMemSetAccess drains almost - # no in-flight DMA, and the reload H2D is already queued on the - # h2d stream by the time the backward graph launches — compute - # and H2D then run concurrently instead of the H2D landing in - # the bubble after the compute segment. The consumer-side wait - # in wait_ready afterwards only installs the compute-stream - # dependency on an already-completed context. - wait_remap_copy_on_stream(reload_context, self.h2d_stream) - for group in active_groups: - group.prepare_remap(reload_context) - runner.local_graph_reload_context = reload_context - runner.local_graph_reload_state = "reload_pending" + runner.local_graph_reload_state = "reload_deferred" return True - def local_graph_backward_finish_prepare(self, runner): - """Compatibility no-op: prepare already handed the full reload to the worker.""" + def local_graph_backward_submit_reload(self, runner): + """Submit the next runner's H2D requests without gating its compute stream.""" + if self._local_graph_ping_pong_enabled(runner): + return self.local_graph_backward_prepare(runner) + if getattr(runner, "local_graph_reload_state", None) == "reload_submitted": + return True + if getattr(runner, "local_graph_reload_state", None) is None: + self.local_graph_backward_prepare(runner) + if getattr(runner, "local_graph_reload_state", None) != "reload_deferred": + return False active_groups = self._local_graph_active_groups(runner) - if not active_groups: + observed = getattr(runner, "local_graph_consumer_slots", ()) + slot_plan = {} + for group_index, slot_index, vmm_address in observed: + if group_index >= len(runner.local_graph_offload_groups): + slot_plan = {} + break + group = runner.local_graph_offload_groups[group_index] + if slot_index >= len(group.allocations) or group.allocations[slot_index].tensor.data_ptr() != vmm_address: + slot_plan = {} + break + slot_plan.setdefault(id(group), set()).add(slot_index) + valid = bool(slot_plan) and all(id(group) in slot_plan for group in active_groups) + if not valid: + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + raise RuntimeError( + f"invalid local graph consumer slots: rank={rank} " + f"runner_id={id(runner)} observed={observed} " + f"active_groups={[group.name for group in active_groups]}" + ) + runner.local_graph_reload_slot_plan = { + id(group): sorted(slot_plan[id(group)]) for group in active_groups + } + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + self._local_graph_debug( + f"reload_submit_plan rank={rank} runner_id={id(runner)} " + f"consumer_slots={observed} active_groups=" + f"{[(group.name, id(group)) for group in active_groups]} " + f"slot_plan={runner.local_graph_reload_slot_plan}" + ) + runner.local_graph_reload_slots = [] + runner.local_graph_reload_slot_map = {} + for group_index, slot_index, vmm_address in observed: + group = runner.local_graph_offload_groups[group_index] + key = (group_index, slot_index) + if key in runner.local_graph_reload_slot_map: + raise RuntimeError( + f"duplicate local graph consumer slot: runner_id={id(runner)} " + f"group_index={group_index} slot_index={slot_index}" + ) + runner.local_graph_reload_slot_map[key] = len(runner.local_graph_reload_slots) + runner.local_graph_reload_slots.append((group_index, slot_index, vmm_address)) + runner.local_graph_reload_contexts = {} + for group_index, slot_index, _ in runner.local_graph_reload_slots: + group = runner.local_graph_offload_groups[group_index] + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + self._local_graph_debug( + f"reload_submit_slot rank={rank} runner_id={id(runner)} " + f"position={runner.local_graph_reload_slot_map[(group_index, slot_index)]} " + f"group_index={group_index} group={group.name} slot={slot_index}" + ) + allocation = group.allocations[slot_index] + # Remap-only submission: the worker performs the VMM transition but + # does NOT submit H2D yet; that happens at replay time so the copy + # and the graph launch are issued together. + reload_context = remap_only_slot_after( + allocation, group.host_tensors[slot_index], self.h2d_stream + ) + group.prepare_remap(reload_context) + runner.local_graph_reload_contexts[(group_index, slot_index)] = reload_context + runner.local_graph_reload_state = "reload_submitted" + return True + + + if not getattr(runner, "local_graph_reload_groups", None): + return False + if self._local_graph_ping_pong_enabled(runner): return False - if getattr(runner, "local_graph_reload_state", None) != "reload_pending": - raise RuntimeError("local graph reload was not prepared before finish") + # Each group is adopted immediately after its own request is consumed; + # there is no batch remap context left to drain here. + runner.local_graph_reload_context = None + runner.local_graph_reload_groups = None + runner.local_graph_reload_event = None + runner.local_graph_reload_state = None return True def local_graph_backward_wait_ready(self, runner, compute_stream=None): - """Install the reload event dependency when this runner is consumed.""" + """Wait only for this runner's consumer slots.""" active_groups = self._local_graph_active_groups(runner) - step_bytes = self._local_graph_group_bytes( - getattr(runner, "local_graph_offload_groups", []) - ) - self._local_graph_h2d_bytes += step_bytes if not active_groups: - self._local_graph_debug(f"h2d bytes={step_bytes} total_h2d={self._local_graph_h2d_bytes}") return False - if getattr(runner, "local_graph_reload_state", None) is None: - self.local_graph_backward_prepare(runner) - if getattr(runner, "local_graph_reload_state", None) != "reload_pending": - raise RuntimeError("local graph runner reload is not pending for backward consumption") if compute_stream is None: compute_stream = torch.cuda.current_stream() if self._local_graph_ping_pong_enabled(runner): - reload_event = runner.local_graph_reload_event - if reload_event is None: + if runner.local_graph_reload_event is None: raise RuntimeError("local graph resident reload has no completion event") - compute_stream.wait_event(reload_event) + compute_stream.wait_event(runner.local_graph_reload_event) for group in active_groups: group.adopt_resident_reload() runner.local_graph_reload_event = None - else: - reload_context = runner.local_graph_reload_context - wait_remap_copy_on_stream(reload_context, compute_stream) - for group in active_groups: - group.adopt_reload_submission() - runner.local_graph_reload_context = None + runner.local_graph_reload_state = None + return True + if wait_remap_slot_on_stream is None or launch_remap_slot_h2d is None: + raise RuntimeError("TE per-slot remap wait API is unavailable") + if getattr(runner, "local_graph_reload_state", None) != "reload_submitted": + self.local_graph_backward_submit_reload(runner) + contexts = getattr(runner, "local_graph_reload_contexts", {}) + reload_slots = getattr(runner, "local_graph_reload_slots", ()) + # Two-position H2D pipeline: runner K's copies gate on the completion + # event of the graph two execution-positions earlier (K+2), never on + # the immediate predecessor. By Backward(N)'s replay moment the gate + # (Backward(N+2)) has long completed, so slot(N)'s H2D overlaps the + # still-running Backward(N+1) instead of serializing behind it. + previous_runner = getattr(runner, "prev_bwd_runner", None) + pre_previous_runner = getattr(previous_runner, "prev_bwd_runner", None) if previous_runner is not None else None + gate_event = ( + getattr(pre_previous_runner, "bwd_graph_replay_complete_event", None) + if pre_previous_runner is not None and not self._local_graph_ping_pong_enabled(runner) + else None + ) + if gate_event is not None: + self.h2d_stream.wait_event(gate_event) + self._local_graph_debug( + f"h2d_gate rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1} " + f"runner_id={id(runner)} gate_runner_id={id(pre_previous_runner)}" + ) + for group_index, slot_index, _ in reload_slots: + group = runner.local_graph_offload_groups[group_index] + context = contexts.get((group_index, slot_index)) + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else -1 + h2d_begin = time.time_ns() + # The remap worker only performs the VMM transition. Submit the + # H2D from this training thread so the memcpy is ordered with the + # replay that consumes this slot. launch_remap_slot_h2d waits for + # the worker's remap completion before issuing musaMemcpyAsync on + # h2d_stream; the returned slot event is then propagated to both + # streams before the allocation is adopted. + launch_remap_slot_h2d(context, 0, self.h2d_stream) + wait_remap_slot_on_stream(context, 0, self.h2d_stream) + wait_remap_slot_on_stream(context, 0, compute_stream) + self._local_graph_debug( + f"h2d_launch rank={rank} runner_id={id(runner)} " + f"position={runner.local_graph_reload_slot_map.get((group_index, slot_index), -1)} " + f"group_index={group_index} group={group.name} slot={slot_index} " + f"h2d_submit_ns={time.time_ns() - h2d_begin}" + ) + group.allocations[slot_index].adopt_async_remap() + for group in active_groups: + group.state = "mapped" runner.local_graph_reload_state = None - self._local_graph_debug(f"h2d bytes={step_bytes} total_h2d={self._local_graph_h2d_bytes}") + runner.local_graph_reload_contexts = None + runner.local_graph_reload_slots = None + runner.local_graph_reload_slot_map = None + runner.local_graph_reload_groups = None return True + def begin_local_graph_consumer_capture(self, runner): + """Collect saved-tensor slots while capturing one backward graph.""" + if self._local_graph_consumer_context is not None: + raise RuntimeError("Nested local graph consumer capture") + self._local_graph_consumer_context = {"runner": runner, "slots": []} + + def end_local_graph_consumer_capture(self, runner, success=True): + """Publish the consumer batch captured by one backward graph.""" + context = self._local_graph_consumer_context + self._local_graph_consumer_context = None + if context is None: + return + if success: + runner.local_graph_consumer_slots = tuple( + (item["group_index"], item["slot_index"], item["vmm_address"]) + for item in context["slots"] + ) + self._local_graph_debug( + f"consumer_batch rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1} " + f"runner_id={id(runner)} slots={runner.local_graph_consumer_slots}" + ) + def local_graph_backward_mark_consumed(self, runner, consumed_event): """Publish when a resident bank may be overwritten by a later reload.""" if not self._local_graph_ping_pong_enabled(runner): @@ -1927,8 +2234,12 @@ def on_group_start_backward(self): if not self.do_offload: return debug_rank(f"--on_group_start_backward {self}") - # Wait for compute to finish before starting reload - self.h2d_stream.wait_stream(torch.cuda.current_stream()) + # Fence H2D submission with a point event from the current compute + # stream, rather than waiting on the stream as a whole. This keeps + # later graph work from becoming an implicit dependency of this reload. + compute_ready = torch.cuda.Event() + compute_ready.record(torch.cuda.current_stream()) + self.h2d_stream.wait_event(compute_ready) self.bulk_reload() @@ -2076,10 +2387,11 @@ def backward(ctx, grad_output): """Record the backward event and wait for the h2d stream on cuda graph stream.""" debug_rank("FineGrainedOffloadingBackwardRecordFunction backward") mgr = PipelineOffloadManager.get_instance() - # This event connects TE's graph stream with the reload stream so - # backward consumers do not race H2D reloads launched outside the graph. + # Record the graph-stream position for offload bookkeeping, but do not + # wait on the whole H2D stream here. Each consumed slot installs its + # own completion-event dependency in local_graph_backward_wait_ready(); + # waiting on h2d_stream globally would also wait for later layers. torch.cuda.current_stream().record_event(mgr.cuda_graph_event) - torch.cuda.current_stream().wait_stream(mgr.h2d_stream) return (grad_output,) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index feb9540d972..d2608665e7c 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -544,6 +544,7 @@ def _link_backward_runners(cls): local_graph_offload_index = 0 for runner in cls.backward_runners: runner.next_bwd_runner = None + runner.prev_bwd_runner = None uses_local_offload = getattr( runner, "_uses_local_graph_activation_offload", lambda: False )() @@ -556,6 +557,7 @@ def _link_backward_runners(cls): local_graph_offload_index += 1 for runner, next_runner in zip(cls.backward_runners, cls.backward_runners[1:]): runner.next_bwd_runner = next_runner + next_runner.prev_bwd_runner = runner # The head of the backward chain is consumed first, yet no earlier # backward replay exists to look ahead from. Mark it so its reload can # be dispatched right after its forward D2H instead of inside its own @@ -875,7 +877,26 @@ def forward(ctx, runner, is_first_microbatch, *inputs): else: runner.fwd_graph.replay() if runner.local_graph_offload_groups: - offload_manager.local_graph_forward_replay(runner) + with torch.cuda.stream(replay_stream): + runner.fwd_graph_replay_complete_event.record(replay_stream) + offload_manager.local_graph_forward_replay( + runner, runner.fwd_graph_replay_complete_event + ) + # The backward-chain head has no predecessor replay to submit its + # reload from. Dispatching here (at the last forward replay) put + # the RemapWorker's muMemSetAccess in flight while the + # ReleaseWorker was still unmapping the forward D2H burst — two + # worker threads contending on the driver at the same moment. + # Default now defers the head's remap submission to its own + # backward, where local_graph_backward_wait_ready() lazy-submits + # it after the forward releases have drained; set + # LOCAL_GRAPH_HEAD_RELOAD_AT_FORWARD=1 to restore the old + # forward-time dispatch. + if ( + runner.is_first_bwd_runner + and int(os.environ.get("LOCAL_GRAPH_HEAD_RELOAD_AT_FORWARD", "0")) + ): + offload_manager.local_graph_backward_submit_reload(runner) return runner.fwd_graph_output_surface @staticmethod @@ -906,33 +927,12 @@ def backward(ctx, *grads): offload_manager = runner._local_graph_offload_manager() replay_stream = runner.stream if runner.use_stream else torch.cuda.current_stream() offload_manager.local_graph_backward_wait_ready(runner, replay_stream) - # Submit the successor's reload BEFORE this runner's backward replay: - # with synchronous remap, prepare() blocks until VMM transition and H2D - # submission are done. Doing it pre-replay means muMemSetAccess runs - # while the GPU is quiet (the predecessor graph just finished), and the - # successor's H2D is already on the h2d stream when this backward graph - # launches — current-layer compute and next-layer reload H2D overlap on - # the device. Resident ping-pong keeps its async one-layer semantics: - # prepare is cheap there, and deferring past replay preserves overlap. if next_runner is not None and next_runner.local_graph_offload_groups: if offload_manager is None: offload_manager = next_runner._local_graph_offload_manager() - if offload_manager._local_graph_ping_pong_enabled(next_runner): - offload_manager.local_graph_backward_prepare(next_runner) - else: - second_runner = next_runner.next_bwd_runner - offload_manager.local_graph_backward_prepare(next_runner) - # Prime the layer after next as well: its synchronous remap now - # runs while only the (already submitted) next-layer H2D is in - # flight, keeping the remap chain ahead without ever parking a - # context behind a running communication/compute segment. - if ( - second_runner is not None - and second_runner.local_graph_offload_groups - and not offload_manager._local_graph_ping_pong_enabled(second_runner) - ): - offload_manager.local_graph_backward_prepare(second_runner) - + # Submit the next runner's H2D now so it overlaps this graph + # replay; the next graph installs waits only for its own slots. + offload_manager.local_graph_backward_submit_reload(next_runner) if runner.use_stream: runner.stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(runner.stream): @@ -1003,6 +1003,7 @@ def __init__( self.fwd_graph = None self.bwd_graph = None + self.fwd_graph_replay_complete_event = torch.cuda.Event() self.bwd_graph_replay_complete_event = torch.cuda.Event() self.local_graph_offload_groups = [] self.next_bwd_runner = None @@ -1596,8 +1597,12 @@ def create_bwd_graph(self): if FREEZE_GC: gc.freeze() + manager = self._local_graph_offload_manager() if self.local_graph_offload_groups else None + if manager is not None: + manager.begin_local_graph_consumer_capture(self) with torch.cuda.graph(self.bwd_graph, pool=self.mempool): + self._sync_against_side_streams(self.bwd_side_streams) grad_inputs = torch.autograd.grad( @@ -1654,6 +1659,9 @@ def create_bwd_graph(self): # Non-GTP path: record after the side-stream join. self.bwd_completion_event.record() + if manager is not None: + manager.end_local_graph_consumer_capture(self) + # Unfreeze GC. if FREEZE_GC: gc.unfreeze() diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index dbadd98efbf..9609f88567c 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -1145,6 +1145,7 @@ def test_local_graph_backward_prepare_finish_wait_pipeline(monkeypatch): assert runner.local_graph_reload_state == "reload_pending" assert mgr.local_graph_backward_finish_prepare(runner) assert runner.local_graph_reload_state == "reload_pending" + expected_reload_event = runner.local_graph_reload_event assert mgr.local_graph_backward_wait_ready(runner, compute_stream) reverse_names = list(reversed(names)) @@ -1159,12 +1160,9 @@ def test_local_graph_backward_prepare_finish_wait_pipeline(monkeypatch): assert [e[1] for e in log if e[0] == "group.prepare_remap"] == reverse_names assert [e[1] for e in log if e[0] == "group.adopt_reload_submission"] == reverse_names assert [e for e in log if e[0] == "runner.wait_remap_copy_on_stream"] == [ - # Synchronous remap: prepare blocks on the h2d stream first, then the - # consumer wait installs the compute-stream dependency on the - # already-completed context. ("runner.wait_remap_copy_on_stream", "h2d"), - ("runner.wait_remap_copy_on_stream", "compute"), ] + assert ("compute.wait_event", expected_reload_event) in log assert runner.local_graph_reload_state is None assert mgr._local_graph_h2d_bytes == sum(g.logical_bytes for g in groups) @@ -1189,11 +1187,12 @@ def test_local_graph_backward_wait_ready_primes_unprefetched_runner(monkeypatch) assert mgr.local_graph_backward_wait_ready(runner, compute_stream) - assert [entry[0] for entry in log] == [ + assert [e[0] for e in log] == [ "runner.remap_and_copy_after", "runner.wait_remap_copy_on_stream", + "reload_ready.record", "group.prepare_remap", - "runner.wait_remap_copy_on_stream", + "compute.wait_event", "group.adopt_reload_submission", ] @@ -1246,6 +1245,8 @@ def test_local_graph_forward_replay_primes_first_backward_runner(monkeypatch): "group.enqueue_d2h", "group.try_release", "runner.wait_remap_copy_on_stream", + "reload_ready.record", + "compute.wait_event", "group.adopt_reload_submission", ] From 2df0656862eb7724e5a2d7e63cd89c738f132efb Mon Sep 17 00:00:00 2001 From: dydy Date: Fri, 18 Sep 2026 11:02:23 +0800 Subject: [PATCH 6/6] fix: rename musa to cuda --- .../fine_grained_activation_offload.py | 38 +++++++++++-------- megatron/core/tensor_parallel/layers.py | 8 +--- megatron/core/transformer/cuda_graphs.py | 2 +- megatron/training/training.py | 5 +-- 4 files changed, 26 insertions(+), 27 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index f77025c94ae..28601c41278 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -451,7 +451,7 @@ def _eligible(self, tensor): return ( torch.is_tensor(tensor) and not isinstance(tensor, torch.nn.Parameter) - and tensor.device.type == "musa" + and tensor.is_cuda and tensor.numel() >= self.min_tensor_size and not _te_do_not_offload(tensor) and not getattr(tensor, "_do_not_offload", False) @@ -501,7 +501,7 @@ def load_tensor(self, tensor): def allocate(self, manager=None, slot_bank=None, group_index=None): """Allocate fixed-address slots, optionally from a resident shared bank.""" - from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + from transformer_engine.pytorch.vmm_activation import CUDAActivationVMMAllocation assert self.state == "discovering" self.slot_bank = slot_bank @@ -519,7 +519,7 @@ def allocate(self, manager=None, slot_bank=None, group_index=None): device, ) else: - allocation = MUSAActivationVMMAllocation(shape, stride, dtype, device) + allocation = CUDAActivationVMMAllocation(shape, stride, dtype, device) self.allocations.append(allocation) allocation.set_slot_id( f"rank={torch.distributed.get_rank() if torch.distributed.is_initialized() else -1}/" @@ -620,7 +620,7 @@ def rebind_resident(self, manager, slot_bank, group_index): """Replace provisional capture allocations with a shared resident bank.""" if self.resident: return - from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + from transformer_engine.pytorch.vmm_activation import CUDAActivationVMMAllocation old_allocations = self.allocations self.allocations = [] @@ -653,7 +653,7 @@ def capture_snapshot(self, runner=None): host_tensor.copy_(device_tensor, non_blocking=False) self.capture_forward_runner = runner # Keep an explicit stream dependency even on backends where the - # synchronous copy currently happens to drain the stream. MUSA graph + # synchronous copy currently happens to drain the stream. CUDA graph # capture and the caller's stream are not assumed to be identical. self.capture_snapshot_event = torch.cuda.Event() self.capture_snapshot_event.record(torch.cuda.current_stream()) @@ -722,9 +722,9 @@ def enqueue_d2h( print( f"[vmm-offload][rank={rank}] submit_d2h group={self.name} " f"state={self.state} allocations={len(self.allocations)} " - f"bytes={self.logical_bytes} compute_stream={getattr(compute_stream, 'musa_stream', None)} " - f"d2h_stream={getattr(d2h_stream, 'musa_stream', None)} " - f"release_stream={getattr(release_stream, 'musa_stream', None)}", + f"bytes={self.logical_bytes} compute_stream={getattr(compute_stream, 'cuda_stream', None)} " + f"d2h_stream={getattr(d2h_stream, 'cuda_stream', None)} " + f"release_stream={getattr(release_stream, 'cuda_stream', None)}", flush=True, ) # Event-gated per-group release: fall back to the old shared-stream @@ -802,7 +802,7 @@ def _wait_release_context(self): status = dict(context.status()) if status["done"] != 1: # The allocation owns the context and provides the blocking - # compatibility API without synchronizing a MUSA stream. + # compatibility API without synchronizing a CUDA stream. pass for allocation in self.allocations: allocation.wait_for_async_release() @@ -1011,12 +1011,12 @@ def local_graph_resident_allocation( self, bank, name, group_index, tensor_index, shape, stride, dtype, device ): """Get or create one descriptor-compatible allocation in a resident bank.""" - from transformer_engine.pytorch.vmm_activation import MUSAActivationVMMAllocation + from transformer_engine.pytorch.vmm_activation import CUDAActivationVMMAllocation key = (name, group_index, tensor_index, tuple(shape), tuple(stride), dtype, device) allocations = self._local_graph_resident_slots[bank] if key not in allocations: - allocations[key] = MUSAActivationVMMAllocation(shape, stride, dtype, device) + allocations[key] = CUDAActivationVMMAllocation(shape, stride, dtype, device) return allocations[key] def local_graph_capture_snapshot(self, runner): @@ -1235,9 +1235,9 @@ def _enqueue_graph_groups_offload(self, runner, active_groups, compute_stream, g f"[vmm-offload][rank={rank}] submit_graph_d2h " f"runner_id={id(runner)} groups={[group.name for group in active_groups]} " f"allocations={len(allocations)} " - f"compute_stream={getattr(compute_stream, 'musa_stream', None)} " - f"d2h_stream={getattr(self.d2h_stream, 'musa_stream', None)} " - f"release_stream={getattr(self._local_graph_release_stream, 'musa_stream', None)}", + f"compute_stream={getattr(compute_stream, 'cuda_stream', None)} " + f"d2h_stream={getattr(self.d2h_stream, 'cuda_stream', None)} " + f"release_stream={getattr(self._local_graph_release_stream, 'cuda_stream', None)}", flush=True, ) @@ -1321,6 +1321,7 @@ def local_graph_backward_prepare(self, runner, block_until_submitted=True): reload_event.record(self.h2d_stream) runner.local_graph_reload_event = reload_event runner.local_graph_reload_context = None + runner.local_graph_reload_state = "reload_pending" else: # Submit in backward consumption order. The last forward-offloaded # group has the latest release dependency, so waiting on it first @@ -1432,6 +1433,11 @@ def local_graph_backward_wait_ready(self, runner, compute_stream=None): if compute_stream is None: compute_stream = torch.cuda.current_stream() if self._local_graph_ping_pong_enabled(runner): + # The backward-chain head has no predecessor replay to look-ahead + # from. Non-ping-pong wait_ready already lazy-submits; do the same + # here so the first resident H2D is issued before the wait. + if runner.local_graph_reload_event is None: + self.local_graph_backward_submit_reload(runner) if runner.local_graph_reload_event is None: raise RuntimeError("local graph resident reload has no completion event") compute_stream.wait_event(runner.local_graph_reload_event) @@ -1472,7 +1478,7 @@ def local_graph_backward_wait_ready(self, runner, compute_stream=None): # The remap worker only performs the VMM transition. Submit the # H2D from this training thread so the memcpy is ordered with the # replay that consumes this slot. launch_remap_slot_h2d waits for - # the worker's remap completion before issuing musaMemcpyAsync on + # the worker's remap completion before issuing cudaMemcpyAsync on # h2d_stream; the returned slot event is then propagated to both # streams before the allocation is adopted. launch_remap_slot_h2d(context, 0, self.h2d_stream) @@ -1992,7 +1998,7 @@ def _can_manage_tensor_for_offload(tensor): return ( not isinstance(tensor, torch.nn.Parameter) and not torch_stray_tensor - and tensor.device.type in ("cuda", "musa") + and tensor.is_cuda ) def tensor_push(self, tensor): diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 8725021c991..0aa7d33fdf8 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -349,13 +349,7 @@ def forward(self, input_): weight = GTPEmbeddingWeight.apply(self.weight) # Get the embeddings. - # MUSA: F.embedding's backward (embedding_dense_backward) lowers to a - # thrust sort/unique that SYNCHRONIZES -- illegal inside graph capture - # ("unique_by_key: musaErrorStreamCaptureUnsupported", kills - # cuda_graph_impl=full_iteration; verified by miniexp/unique_bwd_probe). - # The index-based gather (weight[ids]) is numerically identical and its - # backward captures cleanly, so prefer it on MUSA. - if self.deterministic_mode or torch.musa.is_available(): + if self.deterministic_mode or torch.cuda.is_available(): output_parallel = weight[masked_input] else: # F.embedding currently has a non-deterministic backward function diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index d2608665e7c..25485bdd4ec 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -2821,7 +2821,7 @@ def _get_fp8_enabled(): ) # Keep activation offload disabled during both TE warmup and graph capture. - # The MUSA offload path uses a separate D2H stream, which cannot be left + # The CUDA offload path uses a separate D2H stream, which cannot be left # running when the graph capture stream is finalized. if self.config.fine_grained_activation_offloading: kwargs['pre_warmup_hook'] = off_interface.disable_offload diff --git a/megatron/training/training.py b/megatron/training/training.py index cc7143848e5..57b8efe9390 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1107,9 +1107,8 @@ def pretrain( timers = get_timers() if args.fine_grained_activation_offloading: - if os.getenv("ACCELERATOR_BACKEND", "").lower() != "musa": - from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu - set_ideal_affinity_for_current_gpu() + from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu + set_ideal_affinity_for_current_gpu() if cfg_container.logger.log_progress: