diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 9234e142c6c..bbeee561110 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -545,7 +545,9 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): """ residual = node.layer_state.residual shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None) - output = layer.mlp.combine(output, shared_expert_output) + output = layer.mlp.combine(output) + output = layer.mlp.postprocess(output, shared_expert_output) + mlp_output_with_bias = (output, None) if hasattr(layer, 'cuda_graphs') and layer.cuda_graphs: layer.mlp.cudagraph_tensor_store.clear() diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py index f83275ed9c6..b87124bab1d 100755 --- a/megatron/core/models/mamba/mamba_layer_specs.py +++ b/megatron/core/models/mamba/mamba_layer_specs.py @@ -20,7 +20,11 @@ from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.mlp import MLP, MLPSubmodules from megatron.core.transformer.spec_utils import ModuleSpec -from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.transformer.transformer_layer import ( + MoETransformerLayer, + TransformerLayer, + TransformerLayerSubmodules, +) moe = get_moe_module_spec( use_te=True, @@ -78,8 +82,7 @@ ), ), moe_layer=ModuleSpec( - # TODO (rwaleffe): change this to be an "MoELayer" to work with CudaGraphs? - module=TransformerLayer, + module=MoETransformerLayer, submodules=TransformerLayerSubmodules( pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add ), diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py index 48ea84566d5..ac6e8b5bf40 100644 --- a/megatron/core/ssm/mamba_layer.py +++ b/megatron/core/ssm/mamba_layer.py @@ -16,6 +16,7 @@ from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import GraphableMegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -85,6 +86,13 @@ def __init__( self.mamba_bda = build_module(submodules.mamba_bda) self.bias_dropout_add_exec_handler = torch.enable_grad + def create_mcore_cudagraph_manager(self, config): + """Register the mamba layer for cudagraphs.""" + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + if not self.config.cuda_graph_scope or CudaGraphScope.mamba in self.config.cuda_graph_scope: + self.cudagraph_manager = CudaGraphManager(config) + def mamba_state_shapes_per_request(self) -> Tuple[Tuple[int], Tuple[int]]: """Returns the Mamba conv and ssm states shapes per request.""" return self.mixer.mamba_state_shapes_per_request() diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 7f3a5ab0365..bf00717ab6c 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -472,6 +472,27 @@ def _fork_rng(): _set_all_rng_states(*current_states) +# Global flag that's toggled whenever inside a checkpointing context +IS_CHECKPOINTING = False + + +def _set_checkpointing(): + """Set state to checkpointing enabled.""" + global IS_CHECKPOINTING + IS_CHECKPOINTING = True + + +def _unset_checkpointing(): + """Unset state to checkpointing enabled.""" + global IS_CHECKPOINTING + IS_CHECKPOINTING = False + + +def is_checkpointing(): + """Check if currently in a checkpoint context.""" + return IS_CHECKPOINTING + + class CheckpointFunction(torch.autograd.Function): """Checkpoint Function @@ -484,6 +505,8 @@ class CheckpointFunction(torch.autograd.Function): @staticmethod def forward(ctx, run_function, distribute_saved_activations, *args): """Forward pass.""" + _set_checkpointing() + ctx.run_function = run_function ctx.distribute_saved_activations = distribute_saved_activations @@ -504,6 +527,7 @@ def forward(ctx, run_function, distribute_saved_activations, *args): # Store everything. ctx.save_for_backward(*args) + _unset_checkpointing() return outputs # pylint: disable=missing-function-docstring @@ -515,6 +539,8 @@ def backward(ctx, *args): "Checkpointing is not compatible with .grad(), " "please use .backward() if possible" ) + _set_checkpointing() + inputs = ctx.saved_tensors if ctx.distribute_saved_activations: safely_set_viewless_tensor_data( @@ -539,6 +565,8 @@ def backward(ctx, *args): ) torch.autograd.backward(outputs, args) grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs) + + _unset_checkpointing() return (None, None) + grads @@ -615,6 +643,14 @@ def __init__(self, fp8=False): def checkpoint(self, run_function, *args): """Checkpoint function.""" + + # If in cuda graph warmup, disable checkpointing, as 'discard_output_and_register_recompute' + # may be called in a separate graph warmup. + from megatron.core.transformer.cuda_graphs import is_graph_warmup + + if is_graph_warmup(): + return run_function(*args) + self.run_function = run_function self.rng_states = _get_all_rng_states() @@ -628,11 +664,14 @@ def checkpoint(self, run_function, *args): def _recompute(self, _): """Used as a hook to recompute the output.""" - if self.ctx is None: - # The recomputation has been triggered already. Just return. + from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup + + # The recomputation has been triggered already. Just return. + # Handle cudagraphs, do nothing if currently in graph warmup + if self.ctx is None or is_graph_warmup(): return - if not torch.autograd._is_checkpoint_valid(): + if not torch.autograd._is_checkpoint_valid() and not is_graph_capturing(): raise RuntimeError( "Checkpointing is not compatible with .grad(), " "please use .backward() if possible" @@ -691,6 +730,12 @@ def discard_output_and_register_recompute(self, hook_tensor): in the forward pass and the gradient of the hook_tensor is computed before the recomputed tensors are used. """ + + from megatron.core.transformer.cuda_graphs import is_graph_warmup + + if is_graph_warmup(): + return + # use resize to release the output tensor memory and still keep the metadata in the tensors. # the metadata is still needed for backward for output in self.outputs: diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 1e3e3edc558..a2bd593579d 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -8,15 +8,16 @@ import time from collections import defaultdict from contextlib import nullcontext -from dataclasses import fields, is_dataclass +from copy import deepcopy +from dataclasses import dataclass, fields, is_dataclass from enum import Enum from functools import partial -from itertools import zip_longest +from itertools import chain, zip_longest from math import ceil -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List import torch -from torch.utils._pytree import tree_flatten +from torch.utils._pytree import tree_map from megatron.core import parallel_state from megatron.core.num_microbatches_calculator import get_num_microbatches @@ -24,9 +25,9 @@ CudaRNGStatesTracker, get_all_rng_states, get_cuda_rng_tracker, + is_checkpointing, ) from megatron.core.transformer.enums import CudaGraphScope -from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import ( @@ -39,6 +40,7 @@ try: import transformer_engine as te # pylint: disable=unused-import + from transformer_engine.pytorch.distributed import is_fp8_activation_recompute_enabled from transformer_engine.pytorch.fp8 import FP8GlobalStateManager from transformer_engine.pytorch.graph import ( make_graphed_callables, @@ -48,6 +50,7 @@ from transformer_engine.pytorch.graph import set_capture_end as te_set_capture_end from transformer_engine.pytorch.graph import set_capture_start as te_set_capture_start from transformer_engine.pytorch.module.base import TransformerEngineBaseModule + from transformer_engine.pytorch.utils import make_weak_ref HAVE_TE_GRAPHS = True except: @@ -61,7 +64,7 @@ HAVE_TQDM = False _IS_GRAPH_CAPTURING = False - +_IS_GRAPH_WARMUP = False logger = logging.getLogger(__name__) # Freeze GC during capture. @@ -79,7 +82,6 @@ def is_graph_capturing(): """Query if currently capturing.""" - global _IS_GRAPH_CAPTURING return _IS_GRAPH_CAPTURING @@ -95,6 +97,39 @@ def _set_capture_end(): _IS_GRAPH_CAPTURING = False +def is_graph_warmup(): + """Query if currently warming up for graph capture.""" + return _IS_GRAPH_WARMUP + + +def _set_warmup_start(): + """Set graph warmup has started.""" + global _IS_GRAPH_WARMUP + _IS_GRAPH_WARMUP = True + + +def _set_warmup_end(): + """Set graph warmup has ended.""" + global _IS_GRAPH_WARMUP + + +@dataclass +class CudagraphBufferMetadata: + """ + Metadata saved to tensors during cudagraph capture. This data will be used to determine + during graph captue when a cudagraph can reuse a buffer or directly write its output into + a subsequent's graph's input. + """ + + is_cudagraph_input: bool = False + is_cudagraph_output: bool = False + input_use_count: int = 0 + cudagraph_reuse_ref_count: int = 0 + capture_reuse_count: int = 0 + fwd_cudagraph_buffer: torch.Tensor = None + bwd_cudagraph_buffer: torch.Tensor = None + + class ArgMetadata: """Arg meta.""" @@ -105,9 +140,64 @@ def __init__(self, arg): self.dtype = arg.dtype self.device = arg.device self.value = arg.data_ptr() + self.requires_grad = arg.requires_grad + if hasattr(arg, "cg_buffer_metadata"): + # Its important this is a reference copy + self.cg_buffer_metadata = arg.cg_buffer_metadata else: self.value = arg + def zeros_like(self): + """Reconstruct a tensor with the properties as the meta arg.""" + + assert self.type == torch.Tensor + return torch.zeros( + *self.shape, dtype=self.dtype, device=self.device, requires_grad=self.requires_grad + ) + + +class TensorReusePool: + """ + A pool-like list of tensors that can be reused as input and output buffers during graph capture. + Also maintains strong references to all tensors created by this pool, so that they will never be + freed by the memory allocator. + """ + + """Record strong references to buffers created by the pool so they cannot be deallocated between + graph captures.""" + tensor_strong_refs: list = [] + + """Record the data_ptrs of buffers created by the pool to check when a tensor came was + allocated from this pool. """ + tensor_strong_refs_dataptrs: set = set() + + """Buffers that have been returned to the pool and are available for reuse. """ + pool: list[torch.Tensor] = [] + + def insert(self, tensor: torch.Tensor): + """Return a tensor to the pool reuse.""" + assert self.owns(tensor) + self.pool.append(tensor) + + def owns(self, tensor: torch.Tensor): + """Check if a tensor was created from this pool.""" + return tensor.data_ptr() in self.tensor_strong_refs_dataptrs + + def get(self, meta: ArgMetadata): + """Try to get a buffer from the pool. If a matching tensor is already in the pool, its + assumed to be available and returned. Otherwise, allocate a new buffer.""" + + assert isinstance(meta, ArgMetadata) + # Find first matching buffer in pool + for i, buf in enumerate(self.pool): + if buf.shape == meta.shape and buf.dtype == meta.dtype and buf.device == meta.device: + return self.pool.pop(i) + + out = meta.zeros_like() + self.tensor_strong_refs.append(out) + self.tensor_strong_refs_dataptrs.add(out.data_ptr()) + return out + def _check_supported_type(meta): """Check if arg meta is a supported type for cudagraph input/outputs.""" @@ -127,33 +217,13 @@ def _check_supported_type(meta): float, StaticInferenceContext, DynamicInferenceContext, + ArgMetadata, } assert meta.type in _SUPPORTED_TYPES or is_dataclass( meta.value ), f"Cudagraphs recieved an arg of type {meta.type} which is not supported." -def _determine_if_transformer_decoder_layer(base_module): - """Determine if the given module is a transformer decoder layer.""" - # import modules here to avoid a circular import - from megatron.core.ssm.mamba_layer import MambaLayer - from megatron.core.transformer.transformer_layer import BaseTransformerLayer, TransformerLayer - - is_potential_decoder_layer = isinstance( - base_module, (TransformerLayer, BaseTransformerLayer, MambaLayer) - ) - if not is_potential_decoder_layer: - return False - if isinstance(base_module, TransformerLayer) and not isinstance( - base_module.cross_attention, IdentityOp - ): - # If the layer has a cross attention, it is not a decoder layer - return False - else: - # Otherwise it is a decoder layer - return True - - def _determine_if_first_last_layer_of_this_vp_chunk(base_module): """Determine if the given module is the first/last layer of the PP+VPP chunk it belongs to. Returns a tuple of two booleans indicating if the module is the first/last layer of the chunk. @@ -163,6 +233,9 @@ def _determine_if_first_last_layer_of_this_vp_chunk(base_module): from megatron.core.transformer.transformer_block import get_num_layers_to_build from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + if not hasattr(base_module, "layer_number"): + return True, True + # find all first/last layers of this PP stage first_layer_numbers = [] last_layer_numbers = [] @@ -218,6 +291,10 @@ def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Gen return gen +fwd_buffer_reuse_ref_count = 0 +bwd_buffer_reuse_ref_count = 0 + + class _CudagraphGlobalRecord: """A global datastructure that records of the ordering of all _CudaGraphRunner's first fwd or bwd passes. 'create_cudagraphs' will use this to create @@ -229,13 +306,16 @@ class _CudagraphGlobalRecord: """A record of fwd and bwd graph creation, populated with 'record_fwd_graph' and 'record_bwd_graph.""" - cudagraph_record = [] - cudagraph_inference_record = [] + cudagraph_record: list[tuple] = [] + cudagraph_inference_record: list[tuple] = [] + + """A pool-like data structure to reuse input and output buffers across cudagraph.""" + tensor_reuse_pool = TensorReusePool() @classmethod - def record_fwd_graph(cls, runner, args, kwargs): + def record_fwd_graph(cls, runner, args, kwargs, out): """Record a fwd graph to 'cudagraph_record""" - cls.cudagraph_record.append((runner, "fwd", args, kwargs)) + cls.cudagraph_record.append((runner, "fwd", args, kwargs, out)) @classmethod def record_bwd_graph(cls, runner): @@ -246,7 +326,6 @@ def record_bwd_graph(cls, runner): def create_cudagraphs(cls): """Iterate through 'cudagraph_record' creating graphs in the order in which they were recorded.""" - # Cudagraphs have already been created, check that no cudagraphed modules ran in eager mode if cls.cudagraph_created: assert len(cls.cudagraph_record) == 0, ( @@ -260,8 +339,6 @@ def create_cudagraphs(cls): return # Otherwise, create all the recorded cudagraphs. - logging.getLogger(__name__).info(f"Creating {len(cls.cudagraph_record)} CUDA graphs") - has_te_modules = False if HAVE_TE_GRAPHS: for g in cls.cudagraph_record: @@ -270,21 +347,27 @@ def create_cudagraphs(cls): [isinstance(m, TransformerEngineBaseModule) for m in base_module.modules()] ) - # If graphing only transformer layers with self attention, then apply the following - # transformer layer specific optimizations that reduce memory usage and tensor copies: - # These eventually will become unneccessary with: - # https://github.com/pytorch/pytorch/pull/137318 - # 1. Some inputs to TransformerLayer (e.g. rotary_emb) are the same over all layers - # and only need to be set once. - # 2. Because the next layer consumes the previous layer's hidden states, all fwd - # cudagraphs can alternate reusing the same hidden_state input, output buffer. - # Similarly, bwd graphs can alternate the same output, input grad buffers. - optimize_transformer_layer_graph_buffers = all( - [g[0].reuse_input_output_buffer for g in cls.cudagraph_record] - ) - if optimize_transformer_layer_graph_buffers: - prev_fwd_hidden_state_output = None - prev_bwd_hidden_state_inputgrad = None + progress_bar = enumerate(cls.cudagraph_record) + time_start = time.time() + mem_stats_start = torch.cuda.memory_stats() + + if torch.distributed.get_rank() == 0: + if HAVE_TQDM: + progress_bar = tqdm( + progress_bar, "create cuda graphs", total=len(cls.cudagraph_record) + ) + + logger.info(f"Creating {len(cls.cudagraph_record)} CUDA graphs") + if not HAVE_TE_GRAPHS: + logger.warning( + "Transformer Engine was not detected while capturing training cudagraphs." + "As a result cudagraph memory overhead may significantly increase as " + "Transformer Engine's weak reference feature is used on cudagraph input and " + "output buffers. This allows the memory of input and output buffers to be " + " reclaimed across graphs while remaining valid buffers for when the graph " + "is replayed. For more information see: " + "https://github.com/NVIDIA/TransformerEngine/blob/v2.10/transformer_engine/pytorch/utils.py#L759" # pylint: disable=line-too-long + ) gc.collect() torch.cuda.empty_cache() @@ -293,6 +376,8 @@ def create_cudagraphs(cls): if has_te_modules: te_set_capture_start() + global bwd_buffer_reuse_ref_count, fwd_buffer_reuse_ref_count + def format_mem_bytes(mem_bytes): for power, suffix in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb"), (0, "bytes")]: suffix_bytes = 1024**power @@ -300,58 +385,25 @@ def format_mem_bytes(mem_bytes): return "%.1f %s" % (mem_bytes / suffix_bytes, suffix) return "%d bytes" % mem_bytes - time_start = time.time() - mem_stats_start = torch.cuda.memory_stats() - progress_bar = enumerate(cls.cudagraph_record) - if HAVE_TQDM: - progress_bar = tqdm(progress_bar, "create cuda graphs", total=len(cls.cudagraph_record)) for g_idx, g in progress_bar: + if torch.distributed.get_rank() == 0: + mem_stats = torch.cuda.memory_stats() + progress_str = "create cuda graphs | mem: alloc %s, res %s" % ( + format_mem_bytes(mem_stats["allocated_bytes.all.current"]), + format_mem_bytes(mem_stats["reserved_bytes.all.current"]), + ) + if HAVE_TQDM: + progress_bar.set_description(progress_str) + elif g_idx % 100 == 0 or g_idx == len(cls.cudagraph_record) - 1: + logger.info(f"{g_idx}/{len(cls.cudagraph_record)}. {progress_str}") runner, graph_type = g[0:2] - - mem_stats = torch.cuda.memory_stats() - progress_str = "create cuda graphs | mem: alloc %s, res %s" % ( - format_mem_bytes(mem_stats["allocated_bytes.all.current"]), - format_mem_bytes(mem_stats["reserved_bytes.all.current"]), - ) - if HAVE_TQDM: - progress_bar.set_description(progress_str) - elif g_idx % 100 == 0 or g_idx == len(cls.cudagraph_record) - 1: - logger.info(f"{g_idx}/{len(cls.cudagraph_record)}. {progress_str}") - - if optimize_transformer_layer_graph_buffers: - if graph_type == 'fwd': - args, kwargs = g[2:] - - if not runner.is_first_layer: - kwargs['hidden_states'] = prev_fwd_hidden_state_output - runner.create_fwd_graph(args, kwargs, clone_inputs=False) - - # The output of TransformerLayer is: (hidden_states, None) - # The output of MambaLayer is: (hidden_states,) - # make sure to get the hidden states tensor from the tuple - prev_fwd_hidden_state_output = runner.fwd_graph_outputs[0] - - else: - # In vision models, encoder and decoder transformers have different - # hidden_states shapes. Each has its own first and last layers that - # are noncontiguous. Reset prev_bwd_hidden_state_inputgrad to None at - # each last layer to avoid shape mismatch when transitioning between - # encoder and decoder. - if runner.is_last_layer: - prev_bwd_hidden_state_inputgrad = None - - runner.create_bwd_graph(prev_bwd_hidden_state_inputgrad) - - # The first input grad TransformerLayer is for 'hidden_states' - prev_bwd_hidden_state_inputgrad = runner.static_grad_inputs[0] + if graph_type == 'fwd': + args, kwargs, out = g[2:] + runner.create_fwd_graph(args, kwargs, out, clone_inputs=True) else: - runner, graph_type = g[0:2] - if graph_type == 'fwd': - args, kwargs = g[2:] - runner.create_fwd_graph(args, kwargs) - else: - runner.create_bwd_graph() + assert fwd_buffer_reuse_ref_count == 0 + runner.create_bwd_graph() # Memory usage. time_end = time.time() @@ -367,16 +419,18 @@ def format_mem_bytes(mem_bytes): - mem_stats_start["reserved_bytes.all.current"] ), } - logger.info( - "> built %d cuda graph(s) in %.2f sec, with total memory usage: " - "allocated %s, reserved %s." - % ( - len(cls.cudagraph_record), - capture_stats["time"], - format_mem_bytes(capture_stats["allocated_bytes"]), - format_mem_bytes(capture_stats["reserved_bytes"]), + + if torch.distributed.get_rank() == 0: + logger.info( + "> built %d cuda graph(s) in %.2f sec, with total memory usage: " + "allocated %s, reserved %s." + % ( + len(cls.cudagraph_record), + capture_stats["time"], + format_mem_bytes(capture_stats["allocated_bytes"]), + format_mem_bytes(capture_stats["reserved_bytes"]), + ) ) - ) # Mark cuda graphs as created. for g in cls.cudagraph_record: @@ -392,6 +446,8 @@ def format_mem_bytes(mem_bytes): if has_te_modules: te_set_capture_end() + torch.cuda.set_stream(torch.cuda.default_stream()) + # Return capture time and memory usage. return capture_stats @@ -425,8 +481,7 @@ def delete_cuda_graphs(): runner.bwd_graph_recorded = False runner.fwd_graph = None runner.bwd_graph = None - runner.fwd_mempool = None - runner.bwd_mempool = None + runner.mempool = None # Reset global tracking state _CudagraphGlobalRecord.cudagraph_created = False @@ -438,8 +493,6 @@ def delete_cuda_graphs(): torch.cuda.empty_cache() CudaGraphManager.global_mempool = None - CudaGraphManager.fwd_mempools = None - CudaGraphManager.bwd_mempool = None class _GraphStatus(Enum): @@ -500,38 +553,42 @@ def forward(ctx, runner, is_first_microbatch, *inputs): ), "Fwd cudagraph received a different number of tensors than what it was graphed with!" # Copy new data into fwd graph input buffer + need_copy_inputs = [] for user_input, cudagraph_input in zip(inputs, runner.fwd_graph_input_surface): - if user_input.data_ptr() != cudagraph_input.data_ptr(): + if ( + hasattr(cudagraph_input, "can_skip_replay_copy") + and cudagraph_input.can_skip_replay_copy + ): + need_copy_inputs.append(user_input) + assert user_input.data_ptr() == cudagraph_input.data_ptr() + else: cudagraph_input.copy_(user_input) ctx.runner = runner - if runner.fp8_enabled or runner.fp4_enabled: - for m in runner.base_module.modules(): - if isinstance(m, TransformerEngineBaseModule): - m.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() - m.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + ctx.save_for_backward(*need_copy_inputs) - if is_te_min_version("1.13.0"): - FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(m.fp8_meta) - else: - FP8GlobalStateManager.add_fp8_tensors_to_global_buffer( - m.fp8_meta, fp8_weights=m._get_fp8_params() - ) + if runner.fp8_enabled or runner.fp4_enabled: + if isinstance(FP8GlobalStateManager.get_fp8_recipe(), te.common.recipe.DelayedScaling): + for m in runner.base_module.modules(): + if isinstance(m, TransformerEngineBaseModule): + m.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() + m.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + + if is_te_min_version("1.13.0"): + FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(m.fp8_meta) + else: + FP8GlobalStateManager.add_fp8_tensors_to_global_buffer( + m.fp8_meta, fp8_weights=m._get_fp8_params() + ) - is_first_fp8_module = FP8GlobalStateManager.is_first_fp8_module() - if is_first_fp8_module: + # Note that FP8GlobalStateManager.is_first_fp8_module() is inacccurate as each + # layer may be in its own fp8 context, when the fp8 recipe != delayed_scaling + if runner.is_first_layer and (runner.fp8_param_cache_updated != is_first_microbatch): FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(not is_first_microbatch) - ctx.is_first_fp8_module = is_first_fp8_module + runner.fp8_param_cache_updated = is_first_microbatch runner.fwd_graph.replay() - - # if last transformer layer, return a clone of the cudagraph output buffer, as releasing - # the cudagraph output buffer into the rest of the system may allow it to be corrupted - if runner.is_last_layer: - out = tuple(o.clone().detach() for o in runner.fwd_graph_output_surface) - else: - out = tuple(o.detach() for o in runner.fwd_graph_output_surface) - return out + return runner.fwd_graph_output_surface @staticmethod def backward(ctx, *grads): @@ -548,16 +605,28 @@ def backward(ctx, *grads): runner.static_grad_outputs ), "Bwd cudagraph received a different number of tensors than what it was graphed with!" + need_copy_inputs = list(ctx.saved_tensors) + for cudagraph_input in runner.fwd_graph_input_surface: + if ( + hasattr(cudagraph_input, "can_skip_replay_copy") + and cudagraph_input.can_skip_replay_copy + ): + cudagraph_input.copy_(need_copy_inputs.pop(0)) + # Copy new data into bwd graph input buffer for user_output_grad, cudagraph_output_grad in zip(grads, runner.static_grad_outputs): + if cudagraph_output_grad is None: + continue if user_output_grad.data_ptr() != cudagraph_output_grad.data_ptr(): cudagraph_output_grad.copy_(user_output_grad) runner.bwd_graph.replay() runner.status = _GraphStatus.FWD_READY - # Update FP8/FP4 scale factors if needed - if (runner.fp8_enabled or runner.fp4_enabled) and ctx.is_first_fp8_module: + # Update FP8 scale factors if needed + if runner.fp8_enabled and isinstance( + FP8GlobalStateManager.get_fp8_recipe(), te.common.recipe.DelayedScaling + ): FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) # If using gradient_accumulation_fusion, whenever `main_grad` is calculated @@ -566,18 +635,12 @@ def backward(ctx, *grads): for param, grad_added in runner.groundtruth_grad_added_to_main_grad.items(): param.grad_added_to_main_grad = grad_added - grads, is_dummy_grad = runner.get_input_grads_with_dummy_flags() - if runner.is_first_layer: - output_grads = tuple( - b.clone().detach() if not (b is None or dummy) else b - for dummy, b in zip(is_dummy_grad, grads) - ) - else: - output_grads = tuple( - b.detach() if not (b is None or dummy) else b - for dummy, b in zip(is_dummy_grad, grads) - ) - return None, None, *output_grads + # Replaying the next bwd graph destroys the data held in static_grad_inputs, so clone + # wgrads as autograd may launch the next graph before wgrads are accumulated + dgrads = runner.static_grad_inputs[: runner.num_dgrads] + wgrads = (g.clone() for g in runner.static_grad_inputs[runner.num_dgrads :]) + + return None, None, *dgrads, *wgrads class _CudaGraphRunner(torch.nn.Module): @@ -588,23 +651,20 @@ class _CudaGraphRunner(torch.nn.Module): def __init__( self, base_module: MegatronModule, - fwd_mempool: int, - bwd_mempool: int, + mempool: int, fwd_graph_input_args: List[Any], fwd_graph_input_kwargs: Dict[str, Any], - share_cudagraph_io_buffers=None, + func, + need_backward, ): """Creates a _CudaGraphRunner, which holds a single pair of fwd and bwd cudagraphs, which are not created until this runner records its graph creation into - '_CudagraphGlobalRecord', and 'create_cudagraphs()' is called. share_cudagraph_io_buffers - is a boolean flag to indicate whether to reuse the cudagraph input and output buffers for - transformer layer specific optimizations that reduce memory usage and tensor copies.""" + '_CudagraphGlobalRecord', and 'create_cudagraphs()' is called.""" super().__init__() self.base_module = base_module - self.fwd_mempool = fwd_mempool - self.bwd_mempool = bwd_mempool + self.mempool = mempool self.fwd_graph_input_arg_metas = [ArgMetadata(a) for a in fwd_graph_input_args] self.fwd_graph_input_kwarg_metas = { @@ -624,14 +684,30 @@ def __init__( self.fp8_enabled = False self.fp4_enabled = False self.deallocate_pipeline_outputs = False - self.num_warmup_steps = 2 - if isinstance(self.base_module.config, TransformerConfig): + + self.grad_enabled = need_backward and torch.is_grad_enabled() + self.func = super(MegatronModule, self.base_module).__call__ if func is None else func + self.is_first_layer, self.is_last_layer = _determine_if_first_last_layer_of_this_vp_chunk( + base_module + ) + + # We use this attribute to record the value of 'is_first_microbatch' each fwd cudagraph + # replay so that way we only update the value of this flag in FP8GlobalStateManager when + # it changes which incurs an HtoD sync + if self.is_first_layer: + self.fp8_param_cache_updated = None + + if hasattr(self.base_module, "config") and isinstance( + self.base_module.config, TransformerConfig + ): self.fuse_wgrad_accumulation = self.base_module.config.gradient_accumulation_fusion self.backward_retain_grad = self.base_module.config.cuda_graph_retain_backward_graph - self.fp8_enabled = self.base_module.config.fp8 is not None - self.fp4_enabled = self.base_module.config.fp4 is not None self.deallocate_pipeline_outputs = self.base_module.config.deallocate_pipeline_outputs self.num_warmup_steps = self.base_module.config.cuda_graph_warmup_steps + self.fp8_enabled = self.base_module.config.fp8 is not None + self.fp4_enabled = self.base_module.config.fp4 is not None + self.fp8_runtime_enabled = None + self.fp4_runtime_enabled = None if self.fp8_enabled: self.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() @@ -643,84 +719,91 @@ def __init__( self.fp4_recipe = get_fp4_recipe(self.base_module.config) FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False) - # Decide whether to reuse the input and output buffer, and if so, - # whether this layer is the first layer (which needs an input buffer) - # or the last layer (which needs an output buffer) - - self.is_transformer_decoder_layer = _determine_if_transformer_decoder_layer(base_module) - self.reuse_input_output_buffer = ( - share_cudagraph_io_buffers and self.is_transformer_decoder_layer - ) - if self.reuse_input_output_buffer: - self.is_first_layer, self.is_last_layer = ( - _determine_if_first_last_layer_of_this_vp_chunk(base_module) - ) - else: - self.is_first_layer, self.is_last_layer = True, True - def __str__(self): return "%s; hid %s" % ( self.base_module.__class__.__name__, tuple(self.fwd_graph_input_kwarg_metas["hidden_states"].shape), ) - def get_fp8_context(self): - """Return a new fp8 context in cudagraph mode.""" - from megatron.core.fp8_utils import get_fp8_context # to avoid circular import - - return get_fp8_context(self.base_module.config, self.base_module.layer_number - 1) - - def get_fp4_context(self): - """Return a new fp4 context in cudagraph mode.""" - from megatron.core.fp4_utils import get_fp4_context # to avoid circular import - - return get_fp4_context(self.base_module.config, self.base_module.layer_number - 1) - def get_quantization_context(self): """Return appropriate quantization context (FP8 or FP4) in cudagraph mode.""" - if self.fp8_enabled: - return self.get_fp8_context() - elif self.fp4_enabled: - return self.get_fp4_context() + if self.fp8_runtime_enabled: + from megatron.core.fp8_utils import get_fp8_context # to avoid circular import + + return get_fp8_context(self.base_module.config, self.base_module.layer_number - 1) + elif self.fp4_runtime_enabled: + from megatron.core.fp4_utils import get_fp4_context # to avoid circular import + + return get_fp4_context(self.base_module.config, self.base_module.layer_number - 1) else: return nullcontext() - def create_fwd_graph(self, args, kwargs, clone_inputs=True): + def get_connected_params(self, outputs): + """Iterate through the autograd graph of 'outputs' and returns all parameters connected. + In theory this should return all parameters that return a nonzero wgrad when computing + the backward pass of 'outputs'.""" + # Flatten outputs and start traversal from roots that require gradients + args = (outputs,) if torch.is_tensor(outputs) else outputs + stack = [ + t.grad_fn + for t in self.get_tensors(args, check_types=False) + if t.requires_grad and t.grad_fn + ] + visited, p_ids = set(), set() + + while stack: + if (fn := stack.pop()) not in visited: + visited.add(fn) + # AccumulateGrad nodes (leafs) hold the 'variable' (Parameter) they accumulate into + if hasattr(fn, 'variable'): + p_ids.add(id(fn.variable)) + stack.extend(f for f, _ in fn.next_functions if f) + + # Return module params that were found in the graph, preserving original order + return tuple(p for p in self.base_module.parameters() if id(p) in p_ids) + + def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True): """Create a fwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.""" - # Freeze GC, to speed up capture time ~15-20x. - if FREEZE_GC: - gc.freeze() + global fwd_buffer_reuse_ref_count + + self.args = args + self.kwargs = kwargs + self.outputs = outputs # save grads and other variables that may be affected by graph warmup if self.training and torch.is_grad_enabled(): - save_main_grads = [ - param.main_grad.clone() - for param in self.base_module.parameters() - if hasattr(param, 'main_grad') - ] + grad_backup = [] + for param in self.base_module.parameters(): + grad_backup.append(param.main_grad.clone() if hasattr(param, "main_grad") else None) - saved_fp8_tensors = None + saved_fp8_tensors = None + if self.fp8_enabled: + if is_te_min_version("1.13.0"): + saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe) + else: + saved_fp8_tensors = save_fp8_tensors( + [self.base_module], self.fp8_recipe.amax_history_len + ) + elif self.fp4_enabled: + if is_te_min_version("2.7.0.dev0"): + saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe) + else: + raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.") - if self.fp8_enabled: - if is_te_min_version("1.13.0"): - saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe) - else: - saved_fp8_tensors = save_fp8_tensors( - [self.base_module], self.fp8_recipe.amax_history_len - ) - elif self.fp4_enabled: - if is_te_min_version("2.7.0.dev0"): - saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe) - else: - raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.") + # cache the moe aux loss if needed, which is accumulated inside the forward pass + from megatron.core.transformer.transformer_layer import MoETransformerLayer - if clone_inputs: - args, kwargs = self.zero_out_tensors(args, kwargs) + is_moe = isinstance(self.base_module, MoETransformerLayer) + if is_moe: + from megatron.core.transformer.moe.moe_utils import get_moe_layer_wise_logging_tracker - input_tensors = self.get_tensors(args, kwargs) - self.fwd_graph_input_surface = input_tensors + tuple(self.base_module.parameters()) + tracker = get_moe_layer_wise_logging_tracker() + cached_aux_losses = {} + for name in tracker: + if "values" in tracker[name]: + cached_aux_losses[name] = torch.clone(tracker[name]["values"]) self.fwd_graph = torch.cuda.CUDAGraph() @@ -732,183 +815,339 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): _ensure_generator_state_is_cudagraph_safe(gen) ) - # warmup again as case graph capture mode may execute a different codepath - for _ in range(self.num_warmup_steps): - with self.get_quantization_context(): - outputs = self.base_module.forward(*args, **kwargs) - if self.training and torch.is_grad_enabled(): - if isinstance(outputs, torch.Tensor): - outputs = (outputs,) - outputs = self.get_tensors(outputs) - grad_inputs = torch.autograd.grad( - outputs=tuple(o for o in outputs if o.requires_grad), - inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad), - grad_outputs=tuple( - torch.zeros_like(o) if o.requires_grad else None for o in outputs - ), - only_inputs=True, - allow_unused=True, + def _resolve_input_buffer(ten): + if not isinstance(ten, ArgMetadata): + return ten + + # the input tensor is resued from another cudagraph's input or output + if ( + hasattr(ten, "cg_buffer_metadata") + and ten.cg_buffer_metadata.fwd_cudagraph_buffer is not None + ): + global fwd_buffer_reuse_ref_count + buf = ten.cg_buffer_metadata.fwd_cudagraph_buffer + + assert ( + ten.cg_buffer_metadata.is_cudagraph_input + and buf.cg_buffer_metadata.capture_reuse_count > 0 ) - with self.get_quantization_context(): - torch.cuda.synchronize() - # Register default CUDA generators ourselves (fixed in-place to have normal tensors) - # before capture begins, to avoid inference-tensor state issues during capture. - with torch.inference_mode(mode=False): - for device_idx in range(torch.cuda.device_count()): - default_gen = torch.cuda.default_generators[device_idx] - self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(default_gen) + if ( + ten.cg_buffer_metadata.input_use_count > 1 + and ten.cg_buffer_metadata.input_use_count + == buf.cg_buffer_metadata.capture_reuse_count + ): + can_skip_replay_copy = False + else: + can_skip_replay_copy = True + + buf.cg_buffer_metadata.capture_reuse_count -= 1 + if buf.cg_buffer_metadata.capture_reuse_count == 0: + ten.cg_buffer_metadata.fwd_cudagraph_buffer = None + fwd_buffer_reuse_ref_count -= 1 + else: + # need to provide a fresh buffer from the reuse pool + buf = _CudagraphGlobalRecord.tensor_reuse_pool.get(ten) + can_skip_replay_copy = False + + buf = buf.detach().requires_grad_(ten.requires_grad) + buf.can_skip_replay_copy = can_skip_replay_copy + return buf + + if clone_inputs: + # if a buffer is used for multiple inputs, create it now + for ten in self.get_tensors(args, kwargs): + if ( + hasattr(ten, 'cg_buffer_metadata') + and ten.cg_buffer_metadata.input_use_count > 1 + and ten.cg_buffer_metadata.fwd_cudagraph_buffer is None + ): + buf = _CudagraphGlobalRecord.tensor_reuse_pool.get(ten) + buf.cg_buffer_metadata = deepcopy(ten.cg_buffer_metadata) + buf.cg_buffer_metadata.capture_reuse_count = ( + ten.cg_buffer_metadata.input_use_count ) + ten.cg_buffer_metadata.fwd_cudagraph_buffer = buf + fwd_buffer_reuse_ref_count += 1 - with torch.cuda.graph( - self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" - ): - outputs = self.base_module.forward(*args, **kwargs) + self.fwd_graph_input_args = tree_map(_resolve_input_buffer, args) + self.fwd_graph_input_kwargs = tree_map(_resolve_input_buffer, kwargs) + else: + self.fwd_graph_input_args, self.fwd_graph_input_kwargs = args, kwargs + + self.fwd_graph_input_surface = self.get_tensors( + self.fwd_graph_input_args, self.fwd_graph_input_kwargs + ) + + ctx = torch.no_grad() if not self.grad_enabled else nullcontext() + with ctx: + # 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.zeros_like(ten).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, + ) + _set_warmup_end() + + with self.get_quantization_context(): + torch.cuda.synchronize() + # Register default CUDA generators ourselves (fixed in-place to have normal tensors) + # before capture begins, to avoid inference-tensor state issues during capture. + with torch.inference_mode(mode=False): + for device_idx in range(torch.cuda.device_count()): + default_gen = torch.cuda.default_generators[device_idx] + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(default_gen) + ) + + # Freeze GC, to speed up capture time ~15-20x. + if FREEZE_GC: + gc.freeze() + + 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 + ) + + # Unfreeze GC. + if FREEZE_GC: + gc.unfreeze() + + # 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 + # only on the last layer per-device to avoid slowing down graph creation. + if self.is_last_layer: + gc.collect() # save cudagraph output buffer - if isinstance(outputs, torch.Tensor): - outputs = (outputs,) - self.fwd_graph_outputs = outputs - self.fwd_graph_output_surface = self.get_tensors(outputs) + self.fwd_graph_outputs = fwd_graph_outputs + self.fwd_graph_output_surface = self.get_tensors(fwd_graph_outputs) + + for fwd_graph_out, o in zip( + self.fwd_graph_output_surface, self.get_arg_metas(self.outputs) + ): + assert hasattr(o, "cg_buffer_metadata") and o.cg_buffer_metadata.is_cudagraph_output + + if ( + o.cg_buffer_metadata.is_cudagraph_input + and o.cg_buffer_metadata.fwd_cudagraph_buffer is None + ): + fwd_graph_out.cg_buffer_metadata = deepcopy(o.cg_buffer_metadata) + fwd_graph_out.cg_buffer_metadata.capture_reuse_count = ( + o.cg_buffer_metadata.cudagraph_reuse_ref_count + ) + o.cg_buffer_metadata.fwd_cudagraph_buffer = fwd_graph_out + fwd_buffer_reuse_ref_count += 1 + + # if an input buffer requires a copy, and does not have metadata attached to it at this + # point, it will not be reused after this forward pass, so return it to the pool + for buf in self.fwd_graph_input_surface: + if ( + hasattr(buf, "can_skip_replay_copy") + and not buf.can_skip_replay_copy + and not hasattr(buf, "cg_buffer_metadata") + ): + assert _CudagraphGlobalRecord.tensor_reuse_pool.owns(buf) + _CudagraphGlobalRecord.tensor_reuse_pool.insert(buf) if self.training and torch.is_grad_enabled(): assert ( len(self.fwd_graph_output_surface) > 0 - ), """Tried graphing a moudule that returned no tensors in training mode, - however the graphed module must output at least one tensor, + ), """Tried graphing a module that returned no tensors in training mode, + however the graphed module must output at least one tensor, so that a corresponding backward node may be registered in the autograd graph.""" - # restore cached grads - for param in self.base_module.parameters(): - if hasattr(param, 'main_grad'): - saved_grad = save_main_grads.pop(0) - assert ( - param.main_grad.shape == saved_grad.shape - ), "Error restoring grads while cudagraphing!" - param.main_grad.copy_(saved_grad) + self.params_to_backprop = self.get_connected_params(fwd_graph_outputs) + self.num_wgrads = len(self.params_to_backprop) + self.num_dgrads = len(self.fwd_graph_input_surface) + self.fwd_graph_input_surface = self.fwd_graph_input_surface + self.params_to_backprop - if self.fp8_enabled or self.fp4_enabled: - restore_fp8_tensors([self.base_module], saved_fp8_tensors) - - # Unfreeze GC. - if FREEZE_GC: - gc.unfreeze() + if self.fp8_enabled: + restore_fp8_tensors([self.base_module], saved_fp8_tensors) + # restore cached grads + for main_grad_copy, param in zip(grad_backup, self.base_module.parameters()): + if main_grad_copy is not None: + param.main_grad.copy_(main_grad_copy) - # 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 only on the last layer - # per-device to avoid slowing down graph creation. - if self.is_last_layer: - gc.collect() + if is_moe: + for name in tracker: + tracker[name]["values"].copy_(cached_aux_losses[name]) - def create_bwd_graph(self, static_grad_outputs=None): + def create_bwd_graph(self): """Create a bwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.""" - # Freeze GC, to speed up capture time ~15-20x. - if FREEZE_GC: - gc.freeze() + # unlike 'fwd_buffer_reuse_ref_count', 'bwd_buffer_reuse_ref_count' may not decrement + # to 0 when activation checkpointing is used. See [interaction with recompute]. + global bwd_buffer_reuse_ref_count + assert self.grad_enabled self.bwd_graph = torch.cuda.CUDAGraph() # For cases with multiple active RNG states, e.g. TP. for _, state in get_all_rng_states().items(): self.bwd_graph.register_generator_state(state) - if static_grad_outputs is None: - static_grad_outputs = tuple( - torch.zeros_like(o) if o.requires_grad else None - for o in self.fwd_graph_output_surface - ) - else: - # canoncalize as tuple - if torch.is_tensor(static_grad_outputs): - static_grad_outputs = (static_grad_outputs,) + self.static_grad_outputs = [] + for o in self.get_arg_metas(self.outputs): + out_grad = None + if o.requires_grad: + # TODO: (jiemingz) [interaction with recompute] + # for activation recompute, the fwd pass is rerun in the backward pass and + # the metadata we attach in record_graph_capture is lost. As a result the next + # cudagraph expects the buffer to be provided 'fwd_cudagraph_buffer' but is missing. + # So, we cannot always assume this metadata exists. Consequently, there are extra + # copies between the outputs of the fwd-bwd pass and the bwd pass. + if ( + o.cg_buffer_metadata.is_cudagraph_input + and o.cg_buffer_metadata.bwd_cudagraph_buffer is not None + ): + o.cg_buffer_metadata.bwd_cudagraph_buffer.shape == o.shape - torch.cuda.synchronize() - with torch.cuda.graph( - self.bwd_graph, pool=self.bwd_mempool, capture_error_mode="thread_local" - ): + out_grad = o.cg_buffer_metadata.bwd_cudagraph_buffer + o.cg_buffer_metadata.bwd_cudagraph_buffer = None + out_grad.cg_buffer_metadata.capture_reuse_count -= 1 + bwd_buffer_reuse_ref_count -= 1 + else: + out_grad = _CudagraphGlobalRecord.tensor_reuse_pool.get(o) + out_grad.requires_grad = True + self.static_grad_outputs.append(out_grad) + + # Freeze GC, to speed up capture time ~15-20x. + if FREEZE_GC: + gc.freeze() + + with torch.cuda.graph(self.bwd_graph, pool=self.mempool): grad_inputs = torch.autograd.grad( outputs=tuple(o for o in self.fwd_graph_output_surface if o.requires_grad), inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad), - grad_outputs=tuple(o for o in static_grad_outputs if o is not None), + grad_outputs=tuple(o for o in self.static_grad_outputs if o is not None), retain_graph=self.backward_retain_grad, only_inputs=True, allow_unused=True, ) + # Unfreeze GC. + if FREEZE_GC: + gc.unfreeze() + # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs - # that don't require grad. I couldn't think of a one-liner for this pattern. - static_grad_inputs = [] - grad_idx = 0 - for arg in self.fwd_graph_input_surface: - has_wgrad_fusion = self.fuse_wgrad_accumulation and getattr( - arg, "grad_added_to_main_grad", False - ) - if arg.requires_grad: - if has_wgrad_fusion: - static_grad_inputs.append(None) - else: - static_grad_inputs.append(grad_inputs[grad_idx]) - grad_idx += 1 + # that don't require grad + grad_inputs = list(grad_inputs) + self.static_grad_inputs = [] + for input_tensor in self.get_arg_metas(self.args, self.kwargs): + if input_tensor.requires_grad: + input_grad = grad_inputs.pop(0) + input_grad.cg_buffer_metadata = deepcopy(input_tensor.cg_buffer_metadata) + if input_tensor.cg_buffer_metadata.is_cudagraph_output: + if input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer is None: + input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer = input_grad + input_grad.cg_buffer_metadata.capture_reuse_count += 1 + bwd_buffer_reuse_ref_count += 1 + self.static_grad_inputs.append(input_grad) else: - static_grad_inputs.append(None) + self.static_grad_inputs.append(None) + + # at this point static_grad_inputs hold the input dgrads, add the wgrads next + assert self.num_wgrads == len(grad_inputs) + self.static_grad_inputs.extend(grad_inputs) + self.static_grad_inputs = tuple(self.static_grad_inputs) + self.static_grad_outputs = tuple(self.static_grad_outputs) self.groundtruth_grad_added_to_main_grad = {} if self.fuse_wgrad_accumulation: - for param in self.base_module.parameters(): + for param in self.params_to_backprop: if hasattr(param, "grad_added_to_main_grad"): self.groundtruth_grad_added_to_main_grad[param] = param.grad_added_to_main_grad - self.static_grad_outputs = static_grad_outputs - self.static_grad_inputs = static_grad_inputs + # After backward pass grad_output buffers are no longer used and returned to the pool + for ten in self.static_grad_outputs: + if torch.is_tensor(ten): + # Check that the tensor is not in use. This scenario may occur when a cudagraph + # passes its input directly as an output, and places this output as the + # input of a subsequent cudgraph, leading to a grad output buffer to be still in use + # even after the backward pass. + reuse_count = ( + ten.cg_buffer_metadata.capture_reuse_count + if hasattr(ten, "cg_buffer_metadata") + else 0 + ) - # Unfreeze GC. - if FREEZE_GC: - gc.unfreeze() + if _CudagraphGlobalRecord.tensor_reuse_pool.owns(ten) and reuse_count == 0: + _CudagraphGlobalRecord.tensor_reuse_pool.insert(ten) - if self.is_first_layer: - gc.collect() + # now weakref everything + if HAVE_TE_GRAPHS: - def get_input_grads_with_dummy_flags(self): - """Get the inputs grads that are returned by the bwd cudagraph call. If using grad accum - fusion, wgrads have already been accumulated, so return dummy wgrads.""" + def replace_with_weak_ref(arg): + if not torch.is_tensor(arg): + return arg - is_dummy_grad = [False] * len(self.static_grad_inputs) - if not self.fuse_wgrad_accumulation: - return self.static_grad_inputs, is_dummy_grad - else: - num_dgrads = len(self.static_grad_inputs) - len(list(self.base_module.parameters())) - dgrads = self.static_grad_inputs[:num_dgrads] - wgrads = self.static_grad_inputs[num_dgrads:] - - wgrads_with_placeholders = [] - is_dummy_grad = [False] * len(dgrads) - for idx, param in enumerate(self.base_module.parameters()): - wgrad_is_dummy = getattr(param, "grad_added_to_main_grad", False) - if wgrad_is_dummy: - if getattr(param, "zero_out_wgrad", False): - wgrad = torch.zeros( - param.main_grad.shape, - dtype=param.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - wgrad = torch.empty( - param.main_grad.shape, - dtype=param.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - wgrad = wgrads[idx] - wgrads_with_placeholders.append(wgrad) - is_dummy_grad.append(wgrad_is_dummy) - return tuple(dgrads + wgrads_with_placeholders), is_dummy_grad + ref = make_weak_ref(arg) + ref.requires_grad = arg.requires_grad + if hasattr(arg, "can_skip_replay_copy"): + ref.can_skip_replay_copy = arg.can_skip_replay_copy + return ref + + self.fwd_graph_input_surface = tree_map( + replace_with_weak_ref, self.fwd_graph_input_surface + ) + self.fwd_graph_input_args = tree_map(replace_with_weak_ref, self.fwd_graph_input_args) + self.fwd_graph_input_kwargs = tree_map( + replace_with_weak_ref, self.fwd_graph_input_kwargs + ) + self.fwd_graph_output_surface = tree_map( + replace_with_weak_ref, self.fwd_graph_output_surface + ) + # It is safe to weakref static_grad_inputs as any inuse input grads have a strong ref + # stored in 'bwd_cudagraph_buffer' + self.static_grad_inputs = tree_map(replace_with_weak_ref, self.static_grad_inputs) + self.static_grad_outputs = tree_map(replace_with_weak_ref, self.static_grad_outputs) + + delattr(self, "args") + delattr(self, "kwargs") + delattr(self, "outputs") + + def apply_cudagraph_record_metadata(self, args, kwargs, outputs): + """Attaches graph capture metadata to all passed in tensors.""" + + for t in self.get_tensors(args, kwargs): + if not hasattr(t, "cg_buffer_metadata"): + t.cg_buffer_metadata = CudagraphBufferMetadata() + + t.cg_buffer_metadata.is_cudagraph_input = True + t.cg_buffer_metadata.input_use_count += 1 + + if t.cg_buffer_metadata.is_cudagraph_output: + t.cg_buffer_metadata.cudagraph_reuse_ref_count += 1 + + # mark all outputs, so that the fwd graph we may reuse cudagraph output buffers as inputs + for o in self.get_tensors(outputs): + o.cg_buffer_metadata = CudagraphBufferMetadata() + o.cg_buffer_metadata.is_cudagraph_output = True def record_graph_capture(self, args, kwargs): """Records the data needed to create this runner's forward cudagraph. @@ -916,21 +1155,8 @@ def record_graph_capture(self, args, kwargs): The actual cudagraph will be created when 'create_cudagraphs()` is called. Subsequent passes should replay the graph.""" - if not self.fwd_graph_recorded: - logger.debug(f"Recording forward graph creation...") - if self.is_transformer_decoder_layer and not self.is_first_layer: - # transformer layers hidden_states are already saved as the output of the previous - # layer's cudagraph so avoid saving again - kwargs_copy = dict(kwargs) - kwargs_copy['hidden_states'] = None - _CudagraphGlobalRecord.record_fwd_graph(self, args, kwargs_copy) - else: - _CudagraphGlobalRecord.record_fwd_graph(self, args, kwargs) - - self.fwd_graph_recorded = True - # Run the forward pass as normal in eager mode. - out = super(MegatronModule, self.base_module).__call__(*args, **kwargs) + out = self.func(*args, **kwargs) if type(out) != tuple: out = (out,) @@ -947,9 +1173,38 @@ def record_graph_capture(self, args, kwargs): ] ) - # autograd nodes return inputs as views, so clone the tensor as returning views may cause - # issues, for instance with pipeline parallelism - return tuple(o.clone() if torch.is_tensor(o) else o for o in out) + if not self.fwd_graph_recorded: + logger.debug(f"Recording forward graph creation...") + + self.apply_cudagraph_record_metadata(args, kwargs, out) + + def _replace_with_meta(arg): + return ArgMetadata(arg) if torch.is_tensor(arg) else arg + + m_args = tree_map(_replace_with_meta, args) + m_kwargs = tree_map(_replace_with_meta, kwargs) + m_out = tree_map(_replace_with_meta, out) + _CudagraphGlobalRecord.record_fwd_graph(self, m_args, m_kwargs, m_out) + + if HAVE_TE_GRAPHS: + if FP8GlobalStateManager.is_fp8_enabled(): + # check if the low precision recipe is either fp4 or fp8 + if is_te_min_version("2.7.0.dev0"): + from transformer_engine.common.recipe import NVFP4BlockScaling + + recipe = FP8GlobalStateManager.get_fp8_recipe() + if isinstance(recipe, NVFP4BlockScaling): + self.fp4_runtime_enabled = True + else: + self.fp8_runtime_enabled = True + else: + self.fp8_runtime_enabled = True + + self.fwd_graph_recorded = True + + if len(out) == 1: + return out[0] + return tuple(out) def replay_graph_capture(self, is_first_microbatch, args, kwargs): """Replay the fwd cuda graph with autograd.""" @@ -962,15 +1217,17 @@ def replay_graph_capture(self, is_first_microbatch, args, kwargs): error_msg = "CUDA graph argument mismatch:\n" + "\n".join(mismatch_errors) raise AssertionError(error_msg) - inp_tensors = self.get_tensors(args, kwargs) - func_args = inp_tensors + tuple(self.parameters()) - out = _CudagraphReplayNode.apply(self, is_first_microbatch, *func_args) - out = list(out) + inp_tensors = self.get_tensors(args, kwargs, check_types=False) + if self.grad_enabled: + func_args = inp_tensors + self.params_to_backprop + else: + func_args = inp_tensors - if torch.is_tensor(self.fwd_graph_outputs): - self.fwd_graph_outputs = [self.fwd_graph_outputs] + out = _CudagraphReplayNode.apply(self, is_first_microbatch, *func_args) - return tuple(out.pop(0) if torch.is_tensor(o) else o for o in self.fwd_graph_outputs) + out_iter = iter(self.to_list(out)) + fwd_outputs = self.to_list(self.fwd_graph_outputs) + return tuple(next(out_iter) if torch.is_tensor(o) else o for o in fwd_outputs) def get_mismatch_errors(self, args, kwargs): """Return list of detailed errors for mismatched cudagraph args.""" @@ -1035,66 +1292,55 @@ def check(val, ref, context): return errors - def zero_out_tensors(self, args, kwargs=None): - """Replace all tensors inside arg, kwargs with zeroed copies.""" + def get_arg_metas(self, args, kwargs=None): + """Replaces all passed in tensors with 'ArgMetadata' and returns them as a list.""" + arg_metas = [] - def clone_tensor(ten): - cloned = torch.zeros_like(ten) - cloned.requires_grad = ten.requires_grad - return cloned + def collect(item): + if isinstance(item, ArgMetadata): + arg_metas.append(item) + return item # tree_map expects a return value to rebuild the tree - def process_arg(arg): - _check_supported_type(ArgMetadata(arg)) - if torch.is_tensor(arg): - return clone_tensor(arg) - elif is_dataclass(arg): - for field in fields(arg): - attr = getattr(arg, field.name) - if torch.is_tensor(attr): - setattr(arg, field.name, clone_tensor(attr)) - return arg - - args_replaced = [] - for arg in args: - args_replaced.append(process_arg(arg)) - if kwargs is None: - return args_replaced - - kwargs_replaced = {} - for k, v in kwargs.items(): - kwargs_replaced[k] = process_arg(v) - - return args_replaced, kwargs_replaced + tree_map(collect, args) + if kwargs is not None: + tree_map(collect, kwargs) - @classmethod - def get_tensors(cls, args, kwargs=None): - """Filter and flatten all tensors from args and kwargs.""" + return arg_metas + + def get_tensors(self, args, kwargs=None, check_types=True): + """ + Filter and flatten all tensors from args and kwargs using list comprehensions + and itertools.chain for faster flattening. + """ def extract_tensors(arg): - _check_supported_type(ArgMetadata(arg)) + if check_types: + _check_supported_type(ArgMetadata(arg)) if torch.is_tensor(arg): return [arg] - elif is_dataclass(arg): - tens = [] - for field in fields(arg): - attr = getattr(arg, field.name) - if torch.is_tensor(attr): - tens.append(attr) - return tens - else: - return [] - tens = [] - args, _ = tree_flatten(args) - for a in args: - tens.extend(extract_tensors(a)) + if is_dataclass(arg): + return [ + attr + for field in fields(arg) + if torch.is_tensor(attr := getattr(arg, field.name)) + ] - if kwargs is not None: - kwargs, _ = tree_flatten(kwargs) - for k in kwargs: - tens.extend(extract_tensors(k)) + return [] - return tuple(tens) + if torch.is_tensor(args): + return (args,) + + args_tens = [tensor for arg in args for tensor in extract_tensors(arg)] if args else [] + kwargs_tens = ( + [tensor for val in kwargs.values() for tensor in extract_tensors(val)] if kwargs else [] + ) + + return tuple(chain(args_tens, kwargs_tens)) + + def to_list(self, x): + """Helper function to wrap an input into a list""" + return [x] if torch.is_tensor(x) else list(x) class CudaGraphManager(torch.nn.Module): @@ -1103,17 +1349,8 @@ class CudaGraphManager(torch.nn.Module): """A global mempool for when 'cuda_graph_use_single_mempool' is used.""" global_mempool = None - """Forward pass mempools, used with cudagraph reuse mode.""" - fwd_mempools = None - - """Backward pass mempool, used with cudagraph reuse mode.""" - bwd_mempool = None - def __init__( - self, - config: TransformerConfig, - share_cudagraph_io_buffers: bool = True, - vp_stage: Optional[int] = None, + self, config: TransformerConfig, base_module=None, function_name=None, need_backward=True ): super().__init__() """Creates a CudaGraphManager to manage CUDA graphs for a Megatron module. @@ -1121,14 +1358,21 @@ def __init__( Args: config: TransformerConfig object containing CUDA graph settings for memory pooling, graph retention, gradient accumulation, FP8/FP4, and warmup steps. - share_cudagraph_io_buffers (bool, optional): (DEPRECATED, will be replaced by - config.cuda_graph_share_io_buffers) If None (default) or True, enables - buffer reuse optimizations for transformer and mamba layers. If False, - disables buffer reuse. """ rng_tracker = get_cuda_rng_tracker() - self.share_cudagraph_io_buffers = share_cudagraph_io_buffers - self.vp_stage = vp_stage + self.need_backward = need_backward + + if function_name is not None: + func = getattr(base_module, function_name) + + def wrapped_func(*args, **kwargs): + out = self(base_module, args, kwargs) + return out + + setattr(base_module, function_name, wrapped_func) + else: + func = None + self.func = func # need to delay the import here to avoid a circular import global HAVE_TE_GRAPHS @@ -1144,57 +1388,28 @@ def __init__( ), "RNG tracker does not support cudagraphs!" assert config.cuda_graph_impl == "local", "Option cuda_graph_impl=local not enabled." - assert ( - "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") - or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" - ), ( - "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " - "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." - ) + if torch.cuda.get_device_capability()[0] < 10: + assert ( + "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") + or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" + ), ( + "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " + "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." + ) - self.cudagraph_runners = [] - self.inference_cudagraphs_lookup_table = defaultdict(lambda: None) + self.cudagraph_runners: list[_CudaGraphRunner] = [] + self.inference_cudagraphs_lookup_table: dict = defaultdict(lambda: None) self.is_first_microbatch = False # Without pipeline parallelism, microbatches execute one at a time. # Therefore modules will always execute in the same order, so cudagraphs # can both be reused and share a single mempool. - if parallel_state.get_pipeline_model_parallel_world_size() == 1: - self.reuse_cudagraphs = True - self.use_single_mempool = True - else: - if config.cuda_graph_use_single_mempool: - self.reuse_cudagraphs = False - self.use_single_mempool = True - else: - self.reuse_cudagraphs = True - self.use_single_mempool = False - - # Mempools are static so that multiple cudagraph managers may share the same mempool - if self.use_single_mempool: - if CudaGraphManager.global_mempool is None: - CudaGraphManager.global_mempool = torch.cuda.graph_pool_handle() - else: - # All cudagraphs in the same microbatch use the same mempool. For pipeline parallelism, - # additonally all bwd passes share the same mempool - if CudaGraphManager.fwd_mempools is None: - CudaGraphManager.fwd_mempools = defaultdict( - lambda: defaultdict(torch.cuda.graph_pool_handle) - ) - CudaGraphManager.bwd_mempool = torch.cuda.graph_pool_handle() - - # Cudagraph stream capture requires no operations on the default stream prior to the - # capture, so change to a side stream. - self.stream = torch.cuda.current_stream() - torch.cuda.set_stream(torch.cuda.Stream()) - - def set_is_first_microbatch(self, is_first_microbatch: bool): - """Update the is_first_microbatch flag for weight caching. - - Args: - is_first_microbatch (bool): Whether this is the first microbatch in the step. - """ - self.is_first_microbatch = is_first_microbatch + self.reuse_cudagraphs = parallel_state.get_pipeline_model_parallel_world_size() == 1 + if CudaGraphManager.global_mempool is None: + CudaGraphManager.global_mempool = torch.cuda.graph_pool_handle() + # Cudagraph stream capture requires no operations on the default stream prior to the + # capture, so change to a side stream. + torch.cuda.set_stream(torch.cuda.Stream()) def call_ddp_preforward_hook(self, module): """Call any DDP pre-forward hooks which are used to launch async data parallel @@ -1213,28 +1428,12 @@ def call_ddp_preforward_hook(self, module): def get_cudagraph_runner(self, megatron_module, args, kwargs): '''Returns a valid cudagraph runner for the current forward call. - For single mempool mode, we create a cudagraph for each call, if the module is called - multiple times per step, for instance in the case of pipeline parallelism. The cudagraph corresponding to this call is the first element of 'self.cudagraph_runners'. We iterate through the list by 1 for each call, and the number of calls is equal to the length of 'self.cudagraph_runners'. Otherwise, we assign a mempool per microbatch, which allows cudagraphs to be reused over different microbatches by tracking their respective fwd and bwd passes.''' - if self.use_single_mempool: - fwd_mempool = CudaGraphManager.global_mempool - bwd_mempool = CudaGraphManager.global_mempool - else: - if megatron_module.config.virtual_pipeline_model_parallel_size is not None: - assert ( - self.vp_stage is not None - ), "vp_stage must be passed if virtual pipeline is enabled" - vpp_rank = self.vp_stage - else: - vpp_rank = 0 - fwd_mempool = CudaGraphManager.fwd_mempools[vpp_rank][len(self.cudagraph_runners)] - bwd_mempool = CudaGraphManager.bwd_mempool - if self.reuse_cudagraphs: is_inference_mode = 'inference_context' in kwargs.keys() and kwargs['inference_context'] if is_inference_mode: @@ -1248,15 +1447,20 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs): runner = self.inference_cudagraphs_lookup_table[padded_batch_dimensions] else: # Todo: For training, we could also cache runners based on input shape. - runner = next( - ( - r - for r in self.cudagraph_runners - if r.status == _GraphStatus.FWD_READY + # If autograd is currently disabled, it doesnt matter if a runner was created + # with or without autograd, so just get the first fwd ready runner. + require_grad = self.need_backward and torch.is_grad_enabled() + + def is_valid(r): + return ( + r.status == _GraphStatus.FWD_READY and not r.get_mismatch_errors(args, kwargs) - ), - None, - ) + and (not require_grad or r.grad_enabled) + ) + + # We must choose the first available runner, as the order of + # self.cudagraph_runners corresponds to the capture order. + runner = next((r for r in self.cudagraph_runners if is_valid(r)), None) if runner is None: if _CudagraphGlobalRecord.cudagraph_created: @@ -1268,11 +1472,11 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs): else: runner = _CudaGraphRunner( megatron_module, - fwd_mempool, - bwd_mempool, + CudaGraphManager.global_mempool, args, kwargs, - self.share_cudagraph_io_buffers, + self.func, + self.need_backward, ) self.cudagraph_runners.append(runner) if is_inference_mode: @@ -1292,11 +1496,11 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs): else: runner = _CudaGraphRunner( megatron_module, - fwd_mempool, - bwd_mempool, + CudaGraphManager.global_mempool, args, kwargs, - self.share_cudagraph_io_buffers, + self.func, + self.need_backward, ) self.cudagraph_runners.append(runner) @@ -1312,21 +1516,13 @@ def __call__(self, megatron_module, args, kwargs): kwargs (dict): The keyword args to be passed to the module. """ - # Set the is_first_microbatch flag on the megatron module if it's the first microbatch - if self.is_first_microbatch and hasattr(megatron_module, 'set_is_first_microbatch'): - megatron_module.set_is_first_microbatch() + is_inference_mode = 'inference_context' in kwargs.keys() and kwargs['inference_context'] + is_in_checkpoint_fwd = is_checkpointing() + if HAVE_TE_GRAPHS: + is_in_checkpoint_fwd = is_in_checkpoint_fwd or is_fp8_activation_recompute_enabled() if _CudagraphGlobalRecord.cudagraph_created: if self.training and torch.is_grad_enabled(): - # param.data_ptr() below is used to trigger any hooks that have attached to the - # parameter. Specifically, this is trying to trigger the param sync hook for the - # APEX optimizer, which triggers param syncs by hooking into any param references. - # However cudagraphs disables this, so we workaround by manually referencing - # params here. For more information see: - # https://github.com/NVIDIA/apex/blob/7001836/apex/contrib/optimizers/distributed_fused_adam.py#L885C9 - for param in megatron_module.parameters(): - param.data_ptr() - # Trigger Mcore DDP pre-forward hooks self.call_ddp_preforward_hook(megatron_module) for module in megatron_module.modules(): @@ -1335,7 +1531,7 @@ def __call__(self, megatron_module, args, kwargs): runner = self.get_cudagraph_runner(megatron_module, args, kwargs) out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs) else: - if 'inference_context' in kwargs.keys() and kwargs['inference_context']: + if is_inference_mode: # Inference generation mode creates graphs immediately runner = self.get_cudagraph_runner(megatron_module, args, kwargs) runner.eval() @@ -1343,7 +1539,7 @@ def __call__(self, megatron_module, args, kwargs): if not runner.fwd_graph_recorded: # Reuse graph input-output buffers for inference local_args, local_kwargs = args, kwargs - if runner.reuse_input_output_buffer and not runner.is_first_layer: + if not runner.is_first_layer: # Find previous layer's runner in the global record try: previous_runner = next( @@ -1364,10 +1560,9 @@ def __call__(self, megatron_module, args, kwargs): # No match found for previous layer, continue with no buffer reuse pass - clone_inputs = not ( - runner.reuse_input_output_buffer and not runner.is_first_layer + runner.create_fwd_graph( + local_args, local_kwargs, outputs=None, clone_inputs=runner.is_first_layer ) - runner.create_fwd_graph(local_args, local_kwargs, clone_inputs=clone_inputs) runner.fwd_graph_recorded = True runner.cudagraph_created = True @@ -1378,9 +1573,7 @@ def __call__(self, megatron_module, args, kwargs): # Now replay the graph out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs) - - elif self.training: - # Training mode + elif self.training or is_in_checkpoint_fwd: runner = self.get_cudagraph_runner(megatron_module, args, kwargs) # check if a layer is frozen during training. if not torch.is_grad_enabled(): @@ -1390,13 +1583,17 @@ def __call__(self, megatron_module, args, kwargs): else: # No cudagraphs were found in training mode with grad disabled, so fallback to # eager since autograd is needed to correctly trace the backward graph. - return super(MegatronModule, megatron_module).__call__(*args, **kwargs) + if self.func is not None: + return self.func(*args, **kwargs) + else: + return super(MegatronModule, megatron_module).__call__(*args, **kwargs) + self.is_first_microbatch = False # If forward only, next replay should be a forward pass as well - if self.training and torch.is_grad_enabled(): - runner.status = _GraphStatus.BWD_READY - else: + if is_inference_mode or not torch.is_grad_enabled(): runner.status = _GraphStatus.FWD_READY + else: + runner.status = _GraphStatus.BWD_READY return out diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index fc849da85c8..c30c107e791 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -169,9 +169,12 @@ def __init__(self, config: TransformerConfig, vp_stage: Optional[int] = None): # Enable cuda graphs. if config.cuda_graph_impl == "local": - from megatron.core.transformer.cuda_graphs import CudaGraphManager + if hasattr(self, "create_mcore_cudagraph_manager"): + self.create_mcore_cudagraph_manager(config) + else: + from megatron.core.transformer.cuda_graphs import CudaGraphManager - self.cudagraph_manager = CudaGraphManager(config, vp_stage=vp_stage) + self.cudagraph_manager = CudaGraphManager(config) elif config.cuda_graph_impl == "transformer_engine": # List to store CUDA graphs. A list of `N` CUDA graphs for this layer where N is # the number of microbatches. Multiple CUDA graphs per layer is required to support @@ -336,11 +339,7 @@ def _should_call_te_cudagraph(self, *args, **kwargs): ) def __call__(self, *args, **kwargs): - if self._should_call_local_cudagraph(*args, **kwargs): - # Set the is_first_microbatch flag for weight caching - current_microbatch = getattr(self, 'current_microbatch', 0) - self.cudagraph_manager.set_is_first_microbatch(current_microbatch == 0) return self.cudagraph_manager(self, args, kwargs) elif self._should_call_te_cudagraph(*args, **kwargs): if not self.cuda_graphs: diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 91d443fd9ec..c83e048eef2 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -118,8 +118,11 @@ def __init__( super(MoELayer, self).__init__( config=config, layer_number=layer_number, pg_collection=pg_collection ) + # If using mcore cudagraphs, recompute is handled by transformer_layer.MoETransformerLayer self.moe_layer_recompute = ( - config.recompute_granularity == 'selective' and "moe" in config.recompute_modules + config.recompute_granularity == 'selective' + and "moe" in config.recompute_modules + and config.cuda_graph_impl != 'local' ) self.shared_experts_recompute = ( config.recompute_granularity == 'selective' @@ -206,6 +209,7 @@ def __init__( # Cudagraph tensor store for resuming the forward pass from the end of the cudagraph. self.cudagraph_tensor_store = MoECudaGraphTensorStore() + self.fwd_execution_map = ["route", "expert_compute", "postprocess"] @maybe_skip_or_early_return_by_cudagraph("route") def route(self, hidden_states: torch.Tensor): @@ -291,19 +295,23 @@ def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tenso return output, mlp_bias - def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]): + def combine(self, output: torch.Tensor): """Combines expert outputs via communication and adds shared expert output. This method uses the token dispatcher to combine the outputs from different - experts (e.g., via an All-to-All communication). It then adds the output - from the shared expert if it exists. + experts (e.g., via an All-to-All communication). """ output = self.token_dispatcher.token_combine(output) + return output + + def postprocess(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]): + """Project the output back from latent dimension to hidden dimension after combine + in latent dimension if needed. Combine expert output with shared_experts if needed.""" + output = self.token_dispatcher.combine_postprocess(output) - # Project the output back from latent dimension to hidden dimension after combine - # in latent dimension. if self.config.moe_latent_size: output, _ = self.fc2_latent_proj(output) + if shared_expert_output is not None: output = output + shared_expert_output return output @@ -315,7 +323,7 @@ def router_and_preprocess(self, hidden_states: torch.Tensor): hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map) return hidden_states, probs, residual - def forward(self, hidden_states: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, intermediate_tensors=None): """Forward pass for the MoE layer. The forward pass comprises four main steps: @@ -337,11 +345,16 @@ def forward(self, hidden_states: torch.Tensor): ) # MoE forward: route -> dispatch -> compute -> combine - def custom_forward(hidden_states): + def custom_forward(hidden_states, intermediate_tensors): try: - shared_expert_output = self.shared_experts_compute(hidden_states) - probs, routing_map = self.route(hidden_states) - hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) + if "route" in self.fwd_execution_map: + shared_expert_output = self.shared_experts_compute(hidden_states) + probs, routing_map = self.route(hidden_states) + hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) + + if intermediate_tensors is not None: + return hidden_states, probs, shared_expert_output + except MoECudaGraphPartialCaptureSignal as e: # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator. # It means we should early-return from the MoE layer forward pass. @@ -350,10 +363,28 @@ def custom_forward(hidden_states): # We need to return the intermediate tensors as CUDA graph outputs. return e.get_early_return_outputs(hidden_states, shared_expert_output) - dispatched_input, probs = self.dispatch(hidden_states, probs) - output, mlp_bias = self.routed_experts_compute(dispatched_input, probs) - assert mlp_bias is None, f"mlp_bias is not supported for {type(self.token_dispatcher)}" - output = self.combine(output, shared_expert_output) + if "expert_compute" in self.fwd_execution_map: + if intermediate_tensors is not None: + hidden_states, probs = intermediate_tensors + + dispatched_input, probs = self.dispatch(hidden_states, probs) + output, mlp_bias = self.routed_experts_compute(dispatched_input, probs) + assert ( + mlp_bias is None + ), f"mlp_bias is not supported for {type(self.token_dispatcher)}" + output = self.combine(output) + + if intermediate_tensors is not None: + return output, mlp_bias + + if "postprocess" in self.fwd_execution_map: + if intermediate_tensors is not None: + output, shared_expert_output = intermediate_tensors + + output = self.postprocess(output, shared_expert_output) + + if intermediate_tensors is not None: + return output return output, mlp_bias @@ -365,11 +396,14 @@ def custom_forward(hidden_states): tensor_parallel.random.get_cuda_rng_tracker, parallel_state.get_tensor_model_parallel_group(), hidden_states, + intermediate_tensors, ) else: - outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states) + outputs = tensor_parallel.checkpoint( + custom_forward, False, hidden_states, intermediate_tensors + ) else: - outputs = custom_forward(hidden_states) + outputs = custom_forward(hidden_states, intermediate_tensors) return outputs diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 3bef7d46924..5721f7da4b7 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1232,6 +1232,10 @@ def wrapped_func(moe_layer, *args, **kwargs): Otherwise, we execute the original function and check if we should raise a signal to early return in CUDA graph capture. """ + + if moe_layer.config.cuda_graph_impl != "transformer_engine": + return func(moe_layer, *args, **kwargs) + # The non-cudagraph codepath just calls the original function. if not is_graph_capturing() and moe_layer.cudagraph_tensor_store.is_empty(): return func(moe_layer, *args, **kwargs) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 8322d44d3bb..f2e26c63cf5 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -436,9 +436,9 @@ def __init__( "no_sync": 4, } self.cuda_dtoh_point = "before_permutation_1" - if ( - config.cuda_graph_impl == "transformer_engine" - and CudaGraphScope.moe_preprocess in config.cuda_graph_scope + if config.cuda_graph_impl != "none" and ( + CudaGraphScope.moe_preprocess in config.cuda_graph_scope + or not self.config.cuda_graph_scope ): self.cuda_dtoh_point = "before_ep_alltoall" if MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index cabad4e15d7..77dc81cfd92 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1676,55 +1676,59 @@ def __post_init__(self): raise ValueError("CUDA graphs not supported with CPU offloading.") if self.cuda_graph_impl == "local": - assert not self.cuda_graph_scope or self.cuda_graph_scope == [ - CudaGraphScope.full_iteration - ], ( - "For local cuda graph implementation, the only valid value for " - "cuda_graph_scope is full_iteration, or an empty list to denote layerwise " - "graphs. To use other scopes, use cuda_graph_impl=transformer_engine." - ) + # local impl doesn't currently distinguish between moe_preproocess or moe_router + # so just set both if either is specified. + if ( + CudaGraphScope.moe_router in self.cuda_graph_scope + or CudaGraphScope.moe_preprocess in self.cuda_graph_scope + ): + if CudaGraphScope.moe_router not in self.cuda_graph_scope: + self.cuda_graph_scope.append(CudaGraphScope.moe_router) + if CudaGraphScope.moe_preprocess not in self.cuda_graph_scope: + self.cuda_graph_scope.append(CudaGraphScope.moe_preprocess) + # Check cuda graph scopes if self.cuda_graph_impl == "transformer_engine": assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( "To use full iteration cuda graph, please use " "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine." ) + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + or CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'cuda_graph_scope must not contain both moe and moe_router.' + if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: + assert ( + CudaGraphScope.moe_router in self.cuda_graph_scope + ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' + if self.num_moe_experts is None or self.num_moe_experts <= 1: assert ( CudaGraphScope.moe not in self.cuda_graph_scope - or CudaGraphScope.moe_router not in self.cuda_graph_scope - ), 'cuda_graph_scope must not contain both moe and moe_router.' - if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: - assert ( - CudaGraphScope.moe_router in self.cuda_graph_scope - ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' - if self.num_moe_experts is None or self.num_moe_experts <= 1: + and CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'moe cuda graph is only supported for MoE.' + else: + if self.moe_layer_freq == 1 or ( + isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + ): + assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( + 'mlp cuda graph is only supported for dense layers, ' + 'but not found in the model.' + ) + if ( + self.moe_expert_capacity_factor is None + or not self.moe_pad_expert_input_to_capacity + ): assert ( CudaGraphScope.moe not in self.cuda_graph_scope - and CudaGraphScope.moe_router not in self.cuda_graph_scope - ), 'moe cuda graph is only supported for MoE.' - else: - if self.moe_layer_freq == 1 or ( - isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + ), 'moe cuda graph is only supported with drop-padding MoE.' + if self.moe_token_dispatcher_type == 'alltoall' and ( + self.moe_expert_capacity_factor is not None + or self.moe_router_padding_for_fp8 ): - assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( - 'mlp cuda graph is only supported for dense layers, ' - 'but not found in the model.' + assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( + 'moe_preprocess cuda graph is not supported when there are ' + 'DtoH copies and synchronizations in the preprocess step.' ) - if ( - self.moe_expert_capacity_factor is None - or not self.moe_pad_expert_input_to_capacity - ): - assert ( - CudaGraphScope.moe not in self.cuda_graph_scope - ), 'moe cuda graph is only supported with drop-padding MoE.' - if self.moe_token_dispatcher_type == 'alltoall' and ( - self.moe_expert_capacity_factor is not None - or self.moe_router_padding_for_fp8 - ): - assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( - 'moe_preprocess cuda graph is not supported when there are ' - 'DtoH copies and synchronizations in the preprocess step.' - ) if self.recompute_granularity: if self.recompute_granularity != "selective": @@ -1734,10 +1738,15 @@ def __post_init__(self): else: # The recompute module should be inside or outside of the graph scope. # Recompute module coverring graph scope is not allowed. - if "moe" in self.recompute_modules: + if ( + self.cuda_graph_impl == "transformer_engine" + and "moe" in self.recompute_modules + ): assert ( CudaGraphScope.moe_router not in self.cuda_graph_scope - ), "moe recompute is not supported with moe_router CUDA graph." + ), "moe recompute is not supported with moe_router CUDA graph with: " + "--cuda-graph-impl transformer_engine." + # Graphed recompute module doesn't accept random number. if ( not self.cuda_graph_scope diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 920c3b8fcba..f575794a819 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -268,6 +268,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, vp_stage: Optional[int] = None, ): + self.submodules_config = submodules super().__init__(config=config, vp_stage=vp_stage) if pg_collection is None: @@ -275,7 +276,6 @@ def __init__( self.pg_collection = pg_collection self.tp_group = pg_collection.tp - self.submodules_config = submodules self.layer_number = layer_number + get_transformer_layer_offset( self.config, vp_stage, get_pg_rank(pg_collection.pp) ) @@ -389,6 +389,7 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph(): if ( not self.is_moe_layer or CudaGraphScope.moe_router not in self.config.cuda_graph_scope + or self.config.cuda_graph_impl == "local" ): # Not a MoE layer, or not capturing the router part. return True @@ -462,6 +463,27 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph(): # self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad self.bias_dropout_add_exec_handler = torch.enable_grad + def create_mcore_cudagraph_manager(self, config): + """Register the transformer layer for cudagraphs.""" + + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + # If full scope, just cudagraph the entire layer + if not self.config.cuda_graph_scope: + self.cudagraph_manager = CudaGraphManager(config) + elif ( + CudaGraphScope.attn in self.config.cuda_graph_scope + and self.submodules_config.self_attention != IdentityOp + ): + self.cudagraph_manager = CudaGraphManager(config) + elif ( + CudaGraphScope.mlp in self.config.cuda_graph_scope + and self.submodules_config.mlp != IdentityOp + ): + # Cudagraphing MoE layers are supposed handled by MoeTransforerLayer + assert not self.is_moe_layer + self.cudagraph_manager = CudaGraphManager(config) + @staticmethod def _get_layer_offset(config: TransformerConfig): """ @@ -635,6 +657,23 @@ def _forward_attention( return hidden_states, context + def _forward_pre_mlp_layernorm(self, hidden_states): + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, + ) + + if self.recompute_pre_mlp_layernorm: + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( + self.pre_mlp_layernorm, hidden_states + ) + else: + with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + + return pre_mlp_layernorm_output + def _forward_mlp(self, hidden_states, inference_context=None): """ Perform a forward pass through the feed-forward layer. @@ -646,23 +685,11 @@ def _forward_mlp(self, hidden_states, inference_context=None): output (Tensor): Transformed hidden states of shape [s, b, h]. """ - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) - # Residual connection. residual = hidden_states # Optional Layer norm post the cross-attention. - if self.recompute_pre_mlp_layernorm: - self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: - pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( - self.pre_mlp_layernorm, hidden_states - ) - else: - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: - pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) nvtx_range_push(suffix="mlp") # Potentially chunk the MLP computation during prefill to minimize the peak activation size @@ -714,12 +741,6 @@ def _forward_mlp(self, hidden_states, inference_context=None): self._set_fc2_residual(residual) mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output) - if self.recompute_pre_mlp_layernorm: - # discard the output of the pre-mlp layernorm and register the recompute - # as a gradient hook of mlp_output_with_bias[0] - self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( - mlp_output_with_bias[0] - ) nvtx_range_pop(suffix="mlp") if ( @@ -734,7 +755,7 @@ def _forward_mlp(self, hidden_states, inference_context=None): # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be # recomputed in backward pass. For example, the router path and the shared expert # path. So only register in one path is risky. - for tensor in mlp_output_with_bias[1:]: + for tensor in mlp_output_with_bias: self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) return list(mlp_output_with_bias) + [residual] else: @@ -759,6 +780,13 @@ def _forward_post_mlp(self, mlp_output_with_bias, residual): self.config.inference_fuse_tp_communication ) + if self.recompute_pre_mlp_layernorm: + # discard the output of the pre-mlp layernorm and register the recompute + # as a gradient hook of mlp_output_with_bias[0] + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="mlp_bda") @@ -1055,7 +1083,12 @@ def _te_cuda_graph_replay(self, *args, **kwargs): self.mlp.cudagraph_tensor_store.clear() nvtx_range_pop(suffix="mlp") + # If we early returned, layernorm recompute hooks were attached to the output buffer + # of the cudagraph, so disable the recompute hooks inside _forward_post_mlp + recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm + self.recompute_pre_mlp_layernorm = False output = self._forward_post_mlp(mlp_output_with_bias, residual) + self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: # If EP overlap is enabled, needs to return same outputs as submodule.attn if self.config.overlap_moe_expert_parallel_comm: @@ -1166,6 +1199,7 @@ def __call__(self, *args, **kwargs): kwargs["dynamic_inference_decode_only"] = kwargs[ 'inference_context' ].is_decode_only() + return super().__call__(*args, **kwargs) def get_layer_norm_weights(self): @@ -1175,3 +1209,153 @@ def get_layer_norm_weights(self): List[Tensor]: A list of layernorm weight tensors. """ return + + +class MoETransformerLayer(TransformerLayer): + """ + A Transformer layer specialized for Mixture-of-Experts (MoE) architectures. + + Implements specific functionality to support CUDA graph capture for MoE layers. + Due to the dynamic nature of MoE, capturing the entire layer in a single CUDA graph + can be challenging. This class supports "partial" CUDA graphs by decomposing the + MLP forward pass into router, expert-compute, and post-process stages. + """ + + def __init__(self, *args, **kwargs): + self.is_moe_layer = True + self.use_partial_cudagraphs = False + self.moe_layer_recompute = False + self.token_dispatcher_attrs = {} + + super().__init__(*args, **kwargs) + + 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_scope`. This allows capturing the static parts of the MoE pass + while leaving the expert computation to execute eagerly. + """ + + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + if not self.config.cuda_graph_scope or CudaGraphScope.moe in self.config.cuda_graph_scope: + self.cudagraph_manager = CudaGraphManager(config) + elif ( + CudaGraphScope.moe_router in self.config.cuda_graph_scope + or CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope + ): + # full MoE layer recompute with partial_cudagraphs. If not partial cudagraphs, MoE + # layer recompute is handled by the moe_layer.MoELayer class + self.moe_layer_recompute = ( + self.config.recompute_granularity == 'selective' + and "moe" in self.config.recompute_modules + and self.config.cuda_graph_impl == "local" + ) + + self.use_partial_cudagraphs = True + self.cudagraph_manager_router = CudaGraphManager( + self.config, self, function_name="_forward_mlp_router" + ) + self.cudagraph_manager_postprocess = CudaGraphManager( + self.config, self, function_name="_forward_mlp_postprocess" + ) + + def _forward_mlp_router(self, hidden_states): + """ + Executes the router phase of the MoE block. + + This includes the pre-MLP layernorm and the routing logic. + This method is isolated so it can be captured by `cudagraph_manager_router`. + """ + + residual = hidden_states + self.mlp.fwd_execution_map = "route" + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + router_outputs = self.mlp(pre_mlp_layernorm_output, intermediate_tensors=()) + + for attr_name in self.mlp.token_dispatcher.cudagraph_attrs: + attr = getattr(self.mlp.token_dispatcher, attr_name) + if torch.is_tensor(attr): + if attr_name in self.token_dispatcher_attrs: + self.token_dispatcher_attrs[attr_name].copy_(attr) + else: + self.token_dispatcher_attrs[attr_name] = attr.detach() + + return residual, *router_outputs + + def _forward_mlp_expert_compute(self, hidden_states, probs): + """ + Executes the actual computation of the experts. + + This phase takes the routing information and inputs, dispatches them to the + appropriate experts, and computes the results. In partial graph modes, this + step runs eagerly between the router and postprocess graph replays. + """ + + for name, attr in self.token_dispatcher_attrs.items(): + setattr(self.mlp.token_dispatcher, name, attr) + + self.mlp.fwd_execution_map = "expert_compute" + return self.mlp(None, intermediate_tensors=(hidden_states, probs)) + + def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_bias): + """ + Executes the post-processing phase of the MoE block. + + Handles combining the expert outputs, applying biases, re-registering + activation recomputation hooks if necessary, and performing the final + Bias-Dropout-Add. This method is isolated so it can be captured by cudagraphs. + + """ + + self.mlp.fwd_execution_map = "postprocess" + output = self.mlp(None, intermediate_tensors=(output, shared_expert_output)) + return self._forward_post_mlp((output, mlp_bias), residual) + + def _forward_mlp(self, hidden_states, inference_context=None): + """ + Orchestrates the MLP forward pass, handling partial CUDA graph execution logic. + + If `use_partial_cudagraphs` is True, this method stitches together the + router, expert_compute, and postprocess calls. + """ + + if inference_context is not None: + assert not self.use_partial_cudagraphs, ( + "Partial cudagraphs for MoEs were detected during inference!" + "Please do not use --cuda-graph-scope moe_router moe_preprocess " + "alongside inference." + ) + + def _forward_mlp_partial_cudagraphs(hidden_states, inference_context=None): + residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router( + hidden_states + ) + expert_output, mlp_bias = self._forward_mlp_expert_compute(hidden_states, probs) + return self._forward_mlp_postprocess( + residual, expert_output, shared_expert_output, mlp_bias + ) + + if self.use_partial_cudagraphs: + if self.moe_layer_recompute: + if self.config.fp8 or self.config.fp4: + from megatron.core.extensions.transformer_engine import te_checkpoint + + return te_checkpoint( + _forward_mlp_partial_cudagraphs, + False, + tensor_parallel.random.get_cuda_rng_tracker, + parallel_state.get_tensor_model_parallel_group(), + hidden_states, + ) + else: + return tensor_parallel.checkpoint( + _forward_mlp_partial_cudagraphs, False, hidden_states + ) + else: + return _forward_mlp_partial_cudagraphs(hidden_states) + else: + return super()._forward_mlp(hidden_states) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8173f7c5bdf..ae6b8d762d8 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1339,13 +1339,14 @@ def validate_args(args, defaults={}): ): args.te_rng_tracker = True warn_rank_0("te_rng_tracker is not enabled, enabling it for CUDA graphs.", args.rank) - assert ( - "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") - or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" - ), ( - "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " - "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." - ) + if args.cuda_graph_impl == "transformer_engine": + assert ( + "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") + or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" + ), ( + "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " + "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." + ) if args.cuda_graph_scope == "full" or ( isinstance(args.cuda_graph_scope, list) and "full" in args.cuda_graph_scope ): diff --git a/megatron/training/training.py b/megatron/training/training.py index 6db5c843844..b0dd217593f 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -139,7 +139,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler from megatron.core.transformer.moe import upcycling_utils -from megatron.core.transformer.moe.moe_utils import track_moe_metrics +from megatron.core.transformer.moe.moe_utils import track_moe_metrics, clear_aux_losses_tracker from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper from megatron.core.parallel_state import ( @@ -2905,6 +2905,8 @@ def get_e2e_base_metrics(): timers('interval-time', log_level=0).start(barrier=True) if args.log_energy: energy_monitor.resume() + if args.num_experts is not None: + clear_aux_losses_tracker() # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). # Some of these only happen at specific iterations. Capture updated FLOPs accumulator