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/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index bcc8c14ebb6..28601c41278 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -1,5 +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 @@ -14,6 +16,23 @@ 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 ( + launch_remap_slot_h2d, + release_hooks_after, + 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_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): """Print debug message for a specific rank when DEBUG is enabled.""" @@ -397,6 +416,452 @@ 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.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 per-slot host functions. + self.release_context = None + self.release_contexts = [] + 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 ( + torch.is_tensor(tensor) + and not isinstance(tensor, torch.nn.Parameter) + 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) + ) + + @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 + + 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): + """Allocate fixed-address slots, optionally from a resident shared bank.""" + from transformer_engine.pytorch.vmm_activation import CUDAActivationVMMAllocation + + assert self.state == "discovering" + 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 = 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}/" + 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( + 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 + 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.""" + 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) + 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())} " + 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 _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: + return + from transformer_engine.pytorch.vmm_activation import CUDAActivationVMMAllocation + + 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) + 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: + allocation.close() + + 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) + self.capture_forward_runner = runner + # Keep an explicit stream dependency even on backends where the + # 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()) + 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.""" + 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, release_stream=None, graph_done_event=None + ): + """Queue this group's D2H copies and hand off release to a stream host func. + + 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}") + 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, '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 + # 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): + """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 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: + 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 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 CUDA 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.""" + 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 + 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 a slot request context to this group's bookkeeping.""" + if not self.device_tensors: + return False + 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): + """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}") + 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 + self.state = "mapped" + return True + + class PipelineOffloadManager: """ Singleton manager for coordinating activation offloading across pipeline stages. @@ -449,6 +914,26 @@ 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._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] + # 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 # Do not offload the last X groups so that the reloading won't block the computing stream. @@ -489,6 +974,565 @@ 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 + 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): + 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 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] = CUDAActivationVMMAllocation(shape, stride, dtype, device) + return allocations[key] + + def local_graph_capture_snapshot(self, runner): + """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(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(forward_runner) + + 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: + 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)} " + 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_group_index = 0 + 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.""" + 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_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"consumer_slots={runner.local_graph_capture_metadata['consumer_slots']}" + ) + 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_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, 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] + 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: + # 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) + + 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, '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, + ) + + 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.""" + groups = getattr(runner, "local_graph_offload_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, 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 + 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 + 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 + # makes every earlier dependency a fast path. + # 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 + ) + runner.local_graph_reload_state = "reload_deferred" + return True + + 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) + 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 + # 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): + """Wait only for this runner's consumer slots.""" + active_groups = self._local_graph_active_groups(runner) + if not active_groups: + return False + 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) + for group in active_groups: + group.adopt_resident_reload() + runner.local_graph_reload_event = 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 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) + 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 + 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): + 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.""" debug_rank(f"pushing offload groups to the delayed queue") @@ -954,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 == "cuda" + and tensor.is_cuda ) def tensor_push(self, tensor): @@ -1047,7 +2091,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() @@ -1197,8 +2240,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() @@ -1346,10 +2393,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,) @@ -1360,18 +2408,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 +2471,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/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..0aa7d33fdf8 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -349,7 +349,7 @@ def forward(self, input_): weight = GTPEmbeddingWeight.apply(self.weight) # Get the embeddings. - if self.deterministic_mode: + 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 782fd6bf14f..25485bdd4ec 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, @@ -489,6 +492,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): @@ -527,8 +531,47 @@ 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 + 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 + runner.prev_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 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 + 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 + # 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): """Create recorded CUDA graphs, then remove the saved-tensor observer.""" @@ -562,6 +605,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 +752,17 @@ 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 + runner.is_first_bwd_runner = False + runner.local_graph_forward_runner = runner + # 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 +865,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): @@ -811,6 +876,27 @@ 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: + 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 @@ -835,6 +921,18 @@ 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 + 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) + 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() + # 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): @@ -844,6 +942,10 @@ def backward(ctx, *grads): runner.bwd_graph.replay() 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 @@ -901,7 +1003,15 @@ 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 + 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 @@ -1039,6 +1149,25 @@ def __str__(self): tuple(self.fwd_graph_input_kwarg_metas["hidden_states"].shape), ) + def _uses_local_graph_activation_offload(self): + """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 []) <= LOCAL_GRAPH_ACTIVATION_OFFLOAD_MODULES + ) + + @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 +1384,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 +1454,43 @@ 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() + local_offload_manager.local_graph_capture_snapshot(self) + 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 @@ -1421,12 +1586,23 @@ 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, self.local_graph_forward_runner + ) + # Freeze GC, to speed up capture time ~15-20x. 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( @@ -1483,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() @@ -2641,12 +2820,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 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 - 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 4f9de546161..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 @@ -1044,6 +1050,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). @@ -1892,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.") @@ -2596,33 +2605,30 @@ 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)}." ) - 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 + self.cuda_graph_impl in ("transformer_engine", "full_iteration", "local") + 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/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 1b25dfee1e4..711bee29a82 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1538,6 +1538,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 @@ -1598,26 +1604,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('.') 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..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 @@ -968,3 +968,351 @@ 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" + expected_reload_event = runner.local_graph_reload_event + 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", "h2d"), + ] + 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) + + +def test_local_graph_backward_wait_ready_primes_unprefetched_runner(monkeypatch): + """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) + 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 [e[0] for e in log] == [ + "runner.remap_and_copy_after", + "runner.wait_remap_copy_on_stream", + "reload_ready.record", + "group.prepare_remap", + "compute.wait_event", + "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", + "reload_ready.record", + "compute.wait_event", + "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..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 ( @@ -65,10 +66,85 @@ 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 + # 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 + + 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() @@ -134,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(