diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index e313113a448..7e483722a69 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py @@ -459,7 +459,8 @@ def hook(*unused): if param in self.param_to_bucket_group: assert param.requires_grad - if self.ddp_config.overlap_grad_reduce: + cudagraph_wgrad_ready_event = getattr(param, '_cudagraph_wgrad_ready_event', None) + if self.ddp_config.overlap_grad_reduce and cudagraph_wgrad_ready_event is None: assert ( param.grad is not None ), 'param.grad being None is not safe when overlap_grad_reduce is True' diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 9051fb9f47e..1e88bf8ddfa 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -586,6 +586,18 @@ def start_grad_sync(self, force_all_reduce: Optional[bool] = False): self.grad_reduce_handle is None ), "Should not have multiple communication calls outstanding at once" + # Local CUDA graph replay is asynchronous with respect to the outer + # autograd hooks. Wait before reading, scaling, or reducing gradients + # accumulated by a replay into this bucket. + current_stream = torch.cuda.current_stream() + waited_event_ids = set() + for bucket in self.buckets: + for param in bucket.params_list: + event = getattr(param, "_cudagraph_wgrad_ready_event", None) + if event is not None and id(event) not in waited_event_ids: + current_stream.wait_event(event) + waited_event_ids.add(id(event)) + # Copy accumulated .main_grad into communication buffer before collective if # .main_grad is not in .grad_data already (e.g., because we want to do local # gradient accumulation in a higher precision). diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 99745d98d3d..22322e1b346 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -25,6 +25,7 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.cuda_graphs import annotate_first_last_layer from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -187,6 +188,9 @@ def __init__( raise ValueError("unexpected layer_type") self.layers.append(layer) + if self.config.cuda_graph_impl == "local": + annotate_first_last_layer(self.layers) + # Required for activation recomputation self.num_layers_per_pipeline_rank = len(self.layers) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 0de90c9cde4..9426741eeeb 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -183,47 +183,18 @@ def zeros_like(self): ) -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 alloc_tensor_from_graph_mempool(meta: ArgMetadata): + """Allocates a tensor specified by a ArgMetadata into the graph mempool.""" - 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.""" + torch._C._cuda_beginAllocateCurrentThreadToPool( + torch.cuda.current_device(), CudaGraphManager.global_mempool + ) + out = meta.zeros_like() + out.is_from_global_mempool = True + out.requires_grad_(meta.requires_grad) - 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 + torch._C._cuda_endAllocateToPool(torch.cuda.current_device(), CudaGraphManager.global_mempool) + return out def tree_map(func, tree): @@ -271,33 +242,11 @@ def _check_supported_type(meta): ), f"Cudagraphs received an arg of type {meta.type} which is not supported." -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. - """ - - # import modules here to avoid a circular import - 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 = [] - vp_size = base_module.config.virtual_pipeline_model_parallel_size or 1 - for i in range(vp_size): - # layer numbers are 1-indexed - layer_offset = get_transformer_layer_offset(base_module.config, vp_stage=i) - num_layers_to_build = get_num_layers_to_build(base_module.config, vp_stage=i) - if num_layers_to_build > 0: - first_layer_numbers.append(layer_offset + 1) - last_layer_numbers.append(layer_offset + num_layers_to_build) - return ( - base_module.layer_number in first_layer_numbers, - base_module.layer_number in last_layer_numbers, - ) +def annotate_first_last_layer(layers): + """Annotate the first and last modules in an ordered layer collection.""" + for i, layer in enumerate(layers): + layer.is_first_layer = i == 0 + layer.is_last_layer = i == len(layers) - 1 def _clone_nested_tensors(value: Any) -> Any: @@ -338,6 +287,55 @@ def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Gen return gen +def make_weakref(ten, inplace=True): + """Creates a weak reference to a tensor by creating a tensor that replaces storage with + raw-pointer wrappers that do not hold a storage reference""" + + # Only graph mempool tensors in the graph mempool (e.g. a previous layer's + # output reused as this graph's input) are safe to weak-ref since their memory is + # driver-pinned with stable addresses. Everything else, including, stray tensors + # from dataclass __post_init__ side-effects (e.g. seq_idx created by + # PackedSeqParams.__post_init__ during dataclasses.replace inside the tree_map) must + # retain strong refs, or it will cause a use-after-free on replay that manifests as a + # segfault under memory pressure. + if not ( + HAVE_TE_GRAPHS and torch.is_tensor(ten) and getattr(ten, "is_from_global_mempool", False) + ): + return ten + + try: + wr = make_weak_ref(ten) + if inplace: + ten.data = wr + wr = ten + + except RuntimeError: + # Fallback to keeping a strong reference. There is a known bug where some + # dtypes (e.g. torch.float64) are not mapped to a representation in + # transformer_engine/pytorch/utils.py. + if torch.distributed.get_rank() == 0: + logger.warning( + f"Could not create weak ref for tensor with dtype {arg.dtype}; " + f"keeping strong ref with a potential memory overhead." + ) + + return wr + + +def create_strong_ref(ten: torch.Tensor): + """Create a strong reference to a tensor that keeps memory allocated""" + + ref = ten.detach() + if hasattr(ten, "is_from_global_mempool"): + ref.is_from_global_mempool = ten.is_from_global_mempool + if hasattr(ten, "cg_buffer_metadata"): + ref.cg_buffer_metadata = deepcopy(ten.cg_buffer_metadata) + if hasattr(ten, "can_skip_replay_copy"): + ref.can_skip_replay_copy = ten.can_skip_replay_copy + ref.requires_grad_(ten.requires_grad) + return ref + + fwd_buffer_reuse_ref_count = 0 bwd_buffer_reuse_ref_count = 0 @@ -356,9 +354,6 @@ class _CudagraphGlobalRecord: 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, out): """Record a fwd graph to 'cudagraph_record""" @@ -416,9 +411,6 @@ def create_cudagraphs(cls): "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() - _set_capture_start() if has_te_modules: te_set_capture_start() @@ -426,11 +418,11 @@ def create_cudagraphs(cls): 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 - if mem_bytes >= suffix_bytes: - return "%.1f %s" % (mem_bytes / suffix_bytes, suffix) - return "%d bytes" % mem_bytes + sign, n = ("-", -mem_bytes) if mem_bytes < 0 else ("", mem_bytes) + for p, s in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb")]: + if n >= 1024**p: + return f"{sign}{n / 1024**p:.1f} {s}" + return f"{sign}{n} bytes" for g_idx, g in progress_bar: if torch.distributed.get_rank() == 0: @@ -601,21 +593,16 @@ 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 ( - hasattr(cudagraph_input, "can_skip_replay_copy") - and cudagraph_input.can_skip_replay_copy - ): - need_copy_inputs.append(user_input) + can_skip_replay_copy = getattr( + cudagraph_input, "can_skip_replay_copy", False + ) and getattr(user_input, "can_skip_replay_copy", True) + if can_skip_replay_copy: assert user_input.data_ptr() == cudagraph_input.data_ptr() - else: - if user_input.data_ptr() != cudagraph_input.data_ptr(): - cudagraph_input.copy_(user_input) + elif user_input.data_ptr() != cudagraph_input.data_ptr(): + cudagraph_input.copy_(user_input) ctx.runner = runner - ctx.save_for_backward(*need_copy_inputs) - 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(): @@ -637,6 +624,12 @@ def forward(ctx, runner, is_first_microbatch, *inputs): runner.fp8_param_cache_updated = is_first_microbatch runner.fwd_graph.replay() + + if runner.is_last_layer: + outputs = tuple(torch.clone(t) for t in runner.fwd_graph_output_surface) + for output in outputs: + output.can_skip_replay_copy = False + return outputs return runner.fwd_graph_output_surface @staticmethod @@ -654,14 +647,6 @@ 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: @@ -670,6 +655,9 @@ def backward(ctx, *grads): cudagraph_output_grad.copy_(user_output_grad) runner.bwd_graph.replay() + runner.bwd_graph_replay_complete_event.record(torch.cuda.current_stream()) + for param in runner.params_to_backprop: + param._cudagraph_wgrad_ready_event = runner.bwd_graph_replay_complete_event runner.status = _GraphStatus.FWD_READY # Update FP8 scale factors if needed @@ -678,18 +666,7 @@ def backward(ctx, *grads): ): FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - # If using gradient_accumulation_fusion, whenever `main_grad` is calculated - # the `grad_added_to_main_grad` attribute is expected to set. However when using - # cudagraphs this doesn't occur so we emulate this behavior here. - for param, grad_added in runner.groundtruth_grad_added_to_main_grad.items(): - param.grad_added_to_main_grad = grad_added - - # 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 + return None, None, *runner.static_grad_inputs, *(None,) * len(runner.params_to_backprop) class _CudaGraphRunner(torch.nn.Module): @@ -722,13 +699,13 @@ def __init__( self.fwd_graph = None self.bwd_graph = None + self.bwd_graph_replay_complete_event = torch.cuda.Event() self.fwd_graph_recorded = False self.bwd_graph_recorded = False self.cudagraph_created = False self.status = _GraphStatus.FWD_READY - self.fuse_wgrad_accumulation = False self.backward_retain_grad = False self.fp8_enabled = False self.fp4_enabled = False @@ -739,9 +716,8 @@ def __init__( 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 - ) + self.is_first_layer = getattr(base_module, "is_first_layer", True) + self.is_last_layer = getattr(base_module, "is_last_layer", True) # 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 @@ -752,7 +728,6 @@ def __init__( 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.deallocate_pipeline_outputs = self.base_module.config.deallocate_pipeline_outputs self.num_warmup_steps = self.base_module.config.cuda_graph_warmup_steps @@ -876,6 +851,8 @@ def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True): _ensure_generator_state_is_cudagraph_safe(gen) ) + args_to_clear_buffers = [] + def _resolve_input_buffer(ten): if not isinstance(ten, ArgMetadata): return ten @@ -884,7 +861,6 @@ def _resolve_input_buffer(ten): 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 ( @@ -903,14 +879,12 @@ def _resolve_input_buffer(ten): 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 + args_to_clear_buffers.append(ten) else: - # need to provide a fresh buffer from the reuse pool - buf = _CudagraphGlobalRecord.tensor_reuse_pool.get(ten) + # need to provide a fresh buffer from the pool + buf = alloc_tensor_from_graph_mempool(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 @@ -922,7 +896,7 @@ def _resolve_input_buffer(ten): 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 = alloc_tensor_from_graph_mempool(ten) buf.cg_buffer_metadata = deepcopy(ten.cg_buffer_metadata) buf.cg_buffer_metadata.capture_reuse_count = ( ten.cg_buffer_metadata.input_use_count @@ -1002,24 +976,31 @@ def clone_ten(ten): if self.is_last_layer: gc.collect() + # Deallocate buffers forwarded from previous graphs that are no longer in use + for arg in args_to_clear_buffers: + arg.cg_buffer_metadata.fwd_cudagraph_buffer = None + fwd_buffer_reuse_ref_count -= 1 + # save cudagraph output buffer 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) + self.get_tensors(fwd_graph_outputs), self.get_arg_metas(self.outputs) ): assert hasattr(o, "cg_buffer_metadata") and o.cg_buffer_metadata.is_cudagraph_output + fwd_graph_out.is_from_global_mempool = True + fwd_graph_out.cg_buffer_metadata = deepcopy(o.cg_buffer_metadata) 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 = ( + buf = create_strong_ref(fwd_graph_out) + buf.cg_buffer_metadata.capture_reuse_count = ( o.cg_buffer_metadata.cudagraph_reuse_ref_count ) - o.cg_buffer_metadata.fwd_cudagraph_buffer = fwd_graph_out + o.cg_buffer_metadata.fwd_cudagraph_buffer = buf fwd_buffer_reuse_ref_count += 1 if self.training and torch.is_grad_enabled(): @@ -1029,8 +1010,13 @@ def clone_ten(ten): however the graphed module must output at least one tensor, so that a corresponding backward node may be registered in the autograd graph.""" + self.fwd_graph_input_surface = tree_map(make_weakref, self.fwd_graph_input_surface) + self.fwd_graph_input_args = tree_map(make_weakref, self.fwd_graph_input_args) + self.fwd_graph_input_kwargs = tree_map(make_weakref, self.fwd_graph_input_kwargs) + self.fwd_graph_outputs = tree_map(make_weakref, self.fwd_graph_outputs) + self.fwd_graph_output_surface = tree_map(make_weakref, self.fwd_graph_output_surface) + 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 @@ -1068,6 +1054,7 @@ def create_bwd_graph(self): self.bwd_graph.register_generator_state(state) self.static_grad_outputs = [] + args_to_clear_buffers = [] for o in self.get_arg_metas(self.outputs): out_grad = None if o.requires_grad: @@ -1081,15 +1068,11 @@ def create_bwd_graph(self): 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 - out_grad = o.cg_buffer_metadata.bwd_cudagraph_buffer - o.cg_buffer_metadata.bwd_cudagraph_buffer = None + args_to_clear_buffers.append(o) 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 + out_grad = alloc_tensor_from_graph_mempool(o) self.static_grad_outputs.append(out_grad) # Freeze GC, to speed up capture time ~15-20x. @@ -1105,11 +1088,22 @@ def create_bwd_graph(self): only_inputs=True, allow_unused=True, ) + # Accumulate wgrads directly into main_grad inside the graph + n_act_grads = sum( + 1 for i in self.fwd_graph_input_surface[: self.num_dgrads] if i.requires_grad + ) + for param, wgrad in zip(self.params_to_backprop, grad_inputs[n_act_grads:]): + if wgrad is not None and not getattr(param, 'grad_added_to_main_grad', False): + param.main_grad.add_(wgrad) # Unfreeze GC. if FREEZE_GC: gc.unfreeze() + for arg in args_to_clear_buffers: + arg.cg_buffer_metadata.bwd_cudagraph_buffer = None + bwd_buffer_reuse_ref_count -= 1 + # 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 @@ -1118,103 +1112,29 @@ def create_bwd_graph(self): 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.is_from_global_mempool = True 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 + + if ( + input_tensor.cg_buffer_metadata.is_cudagraph_output + and input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer is None + ): + buf = create_strong_ref(input_grad) + input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer = buf + buf.cg_buffer_metadata.capture_reuse_count += 1 + bwd_buffer_reuse_ref_count += 1 self.static_grad_inputs.append(input_grad) else: 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) + assert len(self.params_to_backprop) == len(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.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 - - # 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 - ) - - if _CudagraphGlobalRecord.tensor_reuse_pool.owns(ten) and reuse_count == 0: - _CudagraphGlobalRecord.tensor_reuse_pool.insert(ten) - - # now weakref everything - if HAVE_TE_GRAPHS: - - def replace_with_weak_ref(arg): - if not torch.is_tensor(arg): - return arg - - try: - ref = make_weak_ref(arg) - except RuntimeError: - # Fallback to keeping a strong reference. There is a known bug where some - # dtypes (e.g. torch.float64) are not mapped to a representation in - # transformer_engine/pytorch/utils.py. - if torch.distributed.get_rank() == 0: - logger.warning( - f"Could not create weak ref for tensor with dtype {arg.dtype}; " - f"keeping strong ref with a potential memory overhead." - ) - return 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 - - # Weak refs replace tensors with raw-pointer wrappers that do not hold a storage - # reference. Only graph mempool tensors in the graph mempool (e.g. a previous layer's - # output reused as this graph's input) are safe to weak-ref since their memory is - # driver-pinned with stable addresses. We identify them as tensors that are not owned - # by the reuse pool and have the attribute `can_skip_replay_copy` set by - # _resolve_input_buffer. Everything else, including reuse-pool buffers, stray tensors - # from dataclass __post_init__ side-effects (e.g. seq_idx created by - # PackedSeqParams.__post_init__ during dataclasses.replace inside the tree_map) must - # retain strong refs, or it will cause a use-after-free on replay that manifests as a - # segfault under memory pressure. - def replace_with_weak_ref_for_input_surface(arg): - if not torch.is_tensor(arg): - return replace_with_weak_ref(arg) - if not _CudagraphGlobalRecord.tensor_reuse_pool.owns(arg) and hasattr( - arg, 'can_skip_replay_copy' - ): - return replace_with_weak_ref(arg) - return arg - - self.fwd_graph_input_surface = tree_map( - replace_with_weak_ref_for_input_surface, 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 - ) - # Outputs can be weakref'd as they are managed by the graph pool - 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) + # 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(make_weakref, self.static_grad_inputs) + self.static_grad_outputs = tree_map(make_weakref, self.static_grad_outputs) delattr(self, "args") delattr(self, "kwargs") @@ -1500,15 +1420,6 @@ def wrapped_func(*args, eager=False, cache_key=None, **kwargs): ), "RNG tracker does not support cudagraphs!" assert config.cuda_graph_impl == "local", "Option cuda_graph_impl=local not enabled." - 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: list[_CudaGraphRunner] = [] self.custom_cudagraphs_lookup_table: dict = defaultdict(lambda: None) self.is_first_microbatch = False diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index d7a1772b61e..490e4d8cc6d 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -461,6 +461,8 @@ def __init__( 'num_global_tokens_per_local_expert', 'reversed_local_input_permutation_mapping', 'routing_map', + 'hidden_shape', + 'probs', ] self.shared_experts = None diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index b6234fe54fe..0415035ffbe 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -21,6 +21,7 @@ from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.transformer.cuda_graphs import annotate_first_last_layer from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -365,6 +366,8 @@ def build_layer(layer_spec, layer_number): for i, layer_spec in enumerate(self.submodules.layer_specs) ] ) + if self.config.cuda_graph_impl == "local": + annotate_first_last_layer(self.layers) # @TODO: add back account_for_embedding_in_pipeline_split (see issue #293) # In pipeline parallelism, we want to add this LN only to the last stage of the pipeline diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..ddd69f2b513 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -18,7 +18,7 @@ from megatron.core.inference.utils import InferenceMode from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.cuda_graphs import is_graph_capturing +from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup, make_weakref from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.mlp import MLP @@ -1489,15 +1489,13 @@ def _forward_mlp_router(self, hidden_states, padding_mask=None): pre_mlp_layernorm_output, intermediate_tensors=(), padding_mask=padding_mask ) - for attr_name in self.mlp.token_dispatcher.cudagraph_attrs: - obj, name = self._resolve_token_dispatcher_attr(attr_name) - attr = getattr(obj, name) - if torch.is_tensor(attr): - cached_attr = self.token_dispatcher_attrs.get(attr_name) - if torch.is_tensor(cached_attr) and not cached_attr.requires_grad: - cached_attr.copy_(attr) - else: - self.token_dispatcher_attrs[attr_name] = attr.detach() + if is_graph_capturing() and not is_graph_warmup(): + for attr_name in self.mlp.token_dispatcher.cudagraph_attrs: + obj, name = self._resolve_token_dispatcher_attr(attr_name) + attr = getattr(obj, name) + if torch.is_tensor(attr): + attr.is_from_global_mempool = True + self.token_dispatcher_attrs[attr_name] = attr return residual, *router_outputs @@ -1530,14 +1528,18 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b """ - # Restore token dispatcher attributes. During graph warmup, the router capture leaves these - # attrs pointing into cudagraph pool memory; restoring them here ensures the postprocess - # graph captures with valid pointers. - self._restore_token_dispatcher_attrs() - self.mlp.fwd_execution_map = "postprocess" output = apply_module(self.mlp)(None, intermediate_tensors=(output, shared_expert_output)) - return self._forward_post_mlp((output, mlp_bias), residual) + out = self._forward_post_mlp((output, mlp_bias), residual) + + if is_graph_capturing() and not is_graph_warmup(): + for attr_name, attr in self.token_dispatcher_attrs.items(): + weak_ref = make_weakref(attr, inplace=False) + self.token_dispatcher_attrs[attr_name] = weak_ref + obj, name = self._resolve_token_dispatcher_attr(attr_name) + setattr(obj, name, weak_ref) + + return out def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None): """ diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index a09f423881e..d6195cc8970 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -967,6 +967,9 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): # Wrap in Float16Module so it accepts fp32_output argument from get_logprobs wrapped_model = Float16Module(transformer_config, model) + # Cudagraph backward capture assumes the model has DDP so create main_grads for params + for param in wrapped_model.parameters(): + param.main_grad = torch.zeros_like(param) # Create test inputs (batch_size=1 required for thd format with sequence packing) batch_size = 1 diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index b9e9c1ae642..61db660d90d 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -644,6 +644,9 @@ def test_llava_cudagraph_is_last_layer_logic(self): # Move model to CUDA self.llava_model.cuda() + # Cudagraph backward capture assumes the model has DDP so create main_grads for params + for param in self.llava_model.parameters(): + param.main_grad = torch.zeros_like(param) set_current_microbatch(self.llava_model.vision_model, 1) set_current_microbatch(self.llava_model.language_model, 1)