diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index 8358e05a612..45e7f66dfe6 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/megatron/core/models/common/model_chunk_schedule_plan.py @@ -11,6 +11,7 @@ from megatron.core.pipeline_parallel.utils import ( AbstractSchedulePlan, NoopScheduleNode, + ScheduleNode, get_comm_stream, get_comp_stream, ) @@ -28,17 +29,17 @@ class ModelChunkState: class TransformerLayerSchedulePlan: - """Schedule the executing plan of the nodes in a transformer/mtp layer. + """Schedule the execution plan for nodes in a transformer or MTP layer. - This class organizes the sub-modules of a transformer/mtp layer, - including attention, post attention, MLP, dispatch, combine and - mtp post process nodes. + This class organizes the submodules of a transformer or MTP layer, including attention, + MLP, MoE dispatch and combine, optional mHC recomputation, and MTP post-processing nodes. layer (TransformerLayerSchedulePlan) ├── attn (TransformerLayerNode): attention -> layernorm -> router -> dispatch preprocess ├── moe_dispatch (TransformerLayerNode): dispatch All2All ├── mlp (TransformerLayerNode): mlp module - ├── moe_combine (TransformerLayerNode): combine All2All + ├── moe_combine (TransformerLayerNode): combine All2All (incl. MLP-side mHC post-processing) + ├── mhc_recompute (ScheduleNode): optional explicit replay before mHC backward └── mtp_post_process (PostProcessNode): mtp post process Note that MTP layer has the same operation and execution order with TransformerLayer regarding @@ -52,6 +53,7 @@ class TransformerLayerSchedulePlan: moe_dispatch = None mlp = None moe_combine = None + mhc_recompute = None mtp_post_process = None def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_args={}): @@ -97,6 +99,9 @@ def release_state(self): if hasattr(self, 'moe_combine') and self.moe_combine is not None: del self.moe_combine self.moe_combine = None + if hasattr(self, 'mhc_recompute') and self.mhc_recompute is not None: + del self.mhc_recompute + self.mhc_recompute = None if hasattr(self, 'mtp_post_process') and self.mtp_post_process is not None: del self.mtp_post_process self.mtp_post_process = None @@ -109,7 +114,7 @@ def release_state(self): def _build_callable_nodes(self, event, comp_stream, comm_stream, extra_args): """ Builds the callable nodes for the transformer/mtp layer: - attn, mlp, moe_dispatch and moe_combine, and mtp_post_process. + attn, mlp, moe_dispatch, moe_combine, and mtp_post_process. """ from megatron.core.models.gpt.fine_grained_callables import ( TransformerLayerNode, @@ -166,6 +171,24 @@ def create_node(stream, module, name): self.moe_dispatch = NoopScheduleNode() self.moe_combine = NoopScheduleNode() + mhc_recompute_manager = extra_args.get("mhc_recompute_manager") + if mhc_recompute_manager is not None and extra_args.get( + "is_last_layer_in_mhc_recompute_group", False + ): + group_index = extra_args["mhc_recompute_group_index"] + # The group counter restarts per module (decoder / mtp), so fold the + # module tag into the NVTX label to keep profiles unambiguous. + module_tag = extra_args.get("mhc_recompute_module_tag", "decoder") + self.mhc_recompute = ScheduleNode( + mhc_recompute_manager.recompute_now, + comp_stream, + event, + name="mhc_recompute", + forward_nvtx_name=f"mhc/recompute/{module_tag}/group_{group_index}/B", + ) + else: + self.mhc_recompute = None + if is_mtp: self.mtp_post_process = create_node( comp_stream, mtp_post_process_module, "mtp_post_process" @@ -237,6 +260,9 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False) When f_layer and b_layer are not None, forward and backward pass are overlapped as follows: comm_stream: combine_bwd | dispatch_fwd->dispatch_bwd | combine_fwd comp_stream: attn_fwd | mlp_bwd->mlp_bwd_dw->mlp_fwd| attn_bwd + MLP-side mHC post-processing runs inside the combine node on the communication stream. + Group recompute runs on the normal compute stream immediately before the node containing + mHC post-processing backward. For MTP, mtp_post_process_fwd is executed after the combine_fwd in the comp_stream, and mtp_post_process_bwd is executed before the combine_bwd in the comp_stream. @@ -254,6 +280,8 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False) if b_layer is not None: b_grad = b_layer.mtp_post_process.backward(b_grad) + if b_layer.mhc_recompute is not None: + b_layer.mhc_recompute.forward() b_grad = b_layer.moe_combine.backward(b_grad) if f_layer is not None: @@ -369,6 +397,7 @@ def __init__( self._model_chunk_state.decoder_input = decoder_input self._model_chunk_state.labels = labels self._model_chunk_state.mtp_hidden_states = None + self._model_chunk_state.mhc_multistream = None self._model_chunk_state.loss_mask = loss_mask self._model_chunk_state.packed_seq_params = packed_seq_params self._model_chunk_state.padding_mask = padding_mask @@ -389,9 +418,11 @@ def __init__( # build layer schedule plan for each layer. # The methods to obtain layers are different for MTP so we need the other build plan for # MTP. Also, this can help annotate MTP layer so that it can know where MTP is. - self._build_layer_schedule_plan(model.decoder, get_comp_stream, get_comm_stream) self._build_layer_schedule_plan( - getattr(model, "mtp", None), get_comp_stream, get_comm_stream + model.decoder, get_comp_stream, get_comm_stream, module_tag="decoder" + ) + self._build_layer_schedule_plan( + getattr(model, "mtp", None), get_comp_stream, get_comm_stream, module_tag="mtp" ) # build post process @@ -400,14 +431,39 @@ def __init__( model, self._model_chunk_state, self._event, get_comp_stream ) - def _build_layer_schedule_plan(self, module, comp_stream, comm_stream): + def _build_layer_schedule_plan(self, module, comp_stream, comm_stream, module_tag): if module is None: return + + from megatron.core.tensor_parallel.random import CheckpointManager + num_layers = len(module.layers) + config = module.config + use_mhc_recompute = ( + module.training + and torch.is_grad_enabled() + and config.enable_hyper_connections + and config.recompute_granularity == "selective" + and "mhc" in config.recompute_modules + ) + group_size = config.mhc_recompute_layer_num or num_layers + mhc_recompute_manager = ( + CheckpointManager() if use_mhc_recompute and num_layers > 0 else None + ) + group_index = 0 + for layer_idx in range(num_layers): + is_group_end = bool( + mhc_recompute_manager is not None + and (layer_idx == num_layers - 1 or (layer_idx + 1) % group_size == 0) + ) extra_args = { "is_first_layer": layer_idx == 0, "is_last_layer": layer_idx == num_layers - 1, + "mhc_recompute_manager": mhc_recompute_manager, + "is_last_layer_in_mhc_recompute_group": is_group_end, + "mhc_recompute_group_index": group_index, + "mhc_recompute_module_tag": module_tag, } layer_plan = TransformerLayerSchedulePlan( module.layers[layer_idx], @@ -419,6 +475,10 @@ def _build_layer_schedule_plan(self, module, comp_stream, comm_stream): ) self._transformer_layers.append(layer_plan) + if is_group_end and layer_idx != num_layers - 1: + group_index += 1 + mhc_recompute_manager = CheckpointManager() + @property def event(self): """Gets the CUDA event for synchronization.""" diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index eb6669b82f6..34533ada4e6 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -21,7 +21,11 @@ MultiTokenPredictionLayer, get_mtp_layer_offset, ) -from megatron.core.transformer.transformer_layer import TransformerLayer, make_viewless_tensor +from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, + make_viewless_tensor, +) from megatron.core.typed_torch import apply_module, copy_signature from megatron.core.utils import internal_api, nvtx_range_pop, nvtx_range_push @@ -103,6 +107,55 @@ def should_free_input(name, is_moe, config, num_local_experts): return free_input_nodes.get(name, False) +def finalize_decoder_layer_output(node, hidden_states): + """Apply the decoder block boundary at whichever node is terminal for the layer. + + The decoder block boundary (mHC output contraction + final layer norm, see + ``TransformerBlock.postprocess_for_layer_schedule``) must run on the last decoder + layer regardless of whether that layer's terminal schedule node is the MoE combine, + the standalone mHC-post, or the dense MLP. Embedding the boundary only in the MoE + closures skips it for mixed patterns whose final layer is dense (for example + ``moe_layer_freq=[1, 0]``), letting an uncontracted ``[s, b, n*h]`` tensor reach GPT + postprocessing without ``learned_output_contract`` or the final layer norm. Factoring + it here keeps the math independent of layer type and mHC-post placement. + + When MTP is enabled the boundary also produces the pre-contraction mHC multi-stream + consumed by the MTP depths. That side output is detached at its producer so MTP reads + a leaf; ``TransformerLayerNode.backward_impl`` reconnects the accumulated gradient when + the scheduler runs this node's backward, exactly as ``residual`` / ``mlp_h_res`` / + ``mlp_hc_h_post`` are bridged. Storing it undetached would let MTP backward traverse the + decoder mHC graph out of schedule order, producing a second-backward error after saved + tensors are freed or bypassing the point where the contracted and MTP branches merge. + + Args: + node: The terminal ``TransformerLayerNode`` for the layer. + hidden_states: The layer output prior to the decoder boundary. + + Returns: + The node output: contracted + normalized ``[s, b, h]`` on the final decoder layer, + otherwise a viewless view of ``hidden_states``. + """ + # Layer nodes exist only for concrete layers; empty decoder chunks use PostProcessNode. + # MTP layers finalize via submodule_mtp_postprocess_forward, not the decoder boundary. + if node.is_mtp or not node.is_last_layer: + return make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + + output, mhc_multistream = node.chunk_state.model.decoder.postprocess_for_layer_schedule( + hidden_states, return_mhc_multistream=True + ) + # postprocess_for_layer_schedule already makes final-layernorm outputs viewless; keep + # this wrapper for the no-layernorm and mHC contraction-only exits. + output = make_viewless_tensor(inp=output, requires_grad=output.requires_grad, keep_graph=True) + # Detach the pre-contraction multi-stream at its producer so MTP reads a leaf and this + # node's backward_impl reconnects the accumulated gradient under scheduler control. + node.chunk_state.mhc_multistream = ( + node.detach(mhc_multistream) if mhc_multistream is not None else None + ) + return output + + class TransformerLayerState: """State shared within a transformer layer. @@ -163,13 +216,15 @@ def forward_impl(self): padding_mask=self.chunk_state.padding_mask, ) - # Saved for later use + # Saved for later use. Keep this as the GPT preprocess output; the decoder layer input may + # be expanded below by TransformerBlock boundary logic when mHC is enabled. self.chunk_state.decoder_input = decoder_input self.chunk_state.rotary_pos_emb = rotary_pos_emb self.chunk_state.rotary_pos_cos = rotary_pos_cos self.chunk_state.rotary_pos_sin = rotary_pos_sin self.chunk_state.sequence_len_offset = sequence_len_offset self.chunk_state.padding_mask = padding_mask + decoder_input = self.gpt_model.decoder.preprocess_for_layer_schedule(decoder_input) return decoder_input @@ -207,13 +262,9 @@ def forward_impl(self, hidden_states): The logits or loss depending on whether labels are provided. """ - empty_decoder = len(self.gpt_model.decoder.layers) == 0 - layer_norm = self.gpt_model.decoder.final_layernorm - if not self.gpt_model.config.mtp_num_layers and empty_decoder and layer_norm: - hidden_states = layer_norm(hidden_states) - hidden_states = make_viewless_tensor( - inp=hidden_states, requires_grad=True, keep_graph=True - ) + if len(self.gpt_model.decoder.layers) == 0: + # Safe for MTP: empty-decoder MTP stages have final_layernorm=None. + hidden_states = self.gpt_model.decoder.postprocess_for_layer_schedule(hidden_states) # Run GPTModel._postprocess loss = self.gpt_model._postprocess( @@ -296,6 +347,10 @@ def __init__( self.detached = tuple() self.before_detached = tuple() self.is_mtp = extra_args.get("is_mtp", False) + self.mhc_recompute_manager = extra_args.get("mhc_recompute_manager") + self.is_last_layer_in_mhc_recompute_group = extra_args.get( + "is_last_layer_in_mhc_recompute_group", False + ) self.post_wgrad_grad_acc_hooks = None # Create flags to indicate first and last layer @@ -470,16 +525,17 @@ def parameters(self): def build_transformer_layer_callables(layer: TransformerLayer): """Create callables for transformer layer nodes. + Divides the transformer layer's operations into a sequence of smaller, independent functions. This decomposition separates computation-heavy tasks (e.g., self-attention, MLP) from communication-heavy tasks (e.g., MoE's All-to-All). - The five callables are: - 1. Attention (computation) - 2. Post-Attention (computation) - 3. MoE Dispatch (communication) - 4. MLP / MoE Experts (computation) - 5. MoE Combine (communication) + The five callable slots are: + 1. Attention and routing preprocess (computation) + 2. MoE Dispatch (communication) + 3. MLP / MoE Experts (computation) + 4. MoE Combine and MLP-side mHC post-processing (communication) + 5. MTP post-processing (computation, MTP layers only) By assigning these functions to different CUDA streams (e.g., a compute stream and a communication stream), the scheduler can overlap their execution, preventing @@ -494,7 +550,6 @@ def build_transformer_layer_callables(layer: TransformerLayer): - forward_funcs: List of callable functions for the layer - backward_dw: Dict of weight gradient functions for the layer """ - is_moe = isinstance(layer.mlp, MoELayer) enable_deepep = ( layer.config.moe_token_dispatcher_type == "flex" @@ -508,6 +563,8 @@ def build_transformer_layer_callables(layer: TransformerLayer): layer.config.moe_token_dispatcher_type == "flex" and layer.config.moe_flex_dispatcher_backend == "ncclep" ) + is_hyper_connection_layer = isinstance(layer, HyperConnectionTransformerLayer) + is_mhc_layer = is_moe and is_hyper_connection_layer def submodule_attn_forward(node: ScheduleNode, hidden_states: torch.Tensor): """ @@ -516,11 +573,19 @@ def submodule_attn_forward(node: ScheduleNode, hidden_states: torch.Tensor): pre mlp layernorm->router->dispatch preprocess """ - if ( + mhc_recompute_manager = getattr(node, "mhc_recompute_manager", None) + is_last_in_mhc_recompute_group = getattr( + node, "is_last_layer_in_mhc_recompute_group", False + ) + if mhc_recompute_manager is not None: + mhc_recompute_manager.is_last_layer_in_recompute_block = is_last_in_mhc_recompute_group + + using_cuda_graph_replay = ( isinstance(layer, GraphableMegatronModule) and hasattr(layer, 'cuda_graphs') and layer.cuda_graphs - ): + ) + if using_cuda_graph_replay: layer.set_te_cuda_graph_backward_dw_wrapper() forward_func = layer._te_cuda_graph_replay else: @@ -533,8 +598,9 @@ def forward_func( rotary_pos_sin: Optional[Tensor] = None, packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[Tensor] = None, + mhc_recompute_manager=None, ): - hidden_states, _ = layer._forward_attention( + attention_kwargs = dict( hidden_states=hidden_states, attention_mask=attention_mask, rotary_pos_emb=rotary_pos_emb, @@ -543,12 +609,29 @@ def forward_func( packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, ) + if is_hyper_connection_layer: + attention_kwargs["mhc_recompute_manager"] = mhc_recompute_manager + hidden_states, _ = layer._forward_attention(**attention_kwargs) if not isinstance(layer.mlp, MoELayer): return hidden_states, None, None, None + if is_mhc_layer: + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post, residual = layer.mlp_hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + else: + mlp_h_res, mlp_hc_h_post = None, None + residual = hidden_states mlp_norm_manager = off_interface(layer.offload_mlp_norm, hidden_states, "mlp_norm") node.layer_state.mlp_norm_manager = mlp_norm_manager - if layer.recompute_pre_mlp_layernorm: - layer.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + checkpoint_pre_mlp_layernorm = layer.recompute_pre_mlp_layernorm or ( + mhc_recompute_manager is not None and layer.mhc_checkpoint_pre_mlp_layernorm + ) + if checkpoint_pre_mlp_layernorm: + layer.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=mhc_recompute_manager + ) with mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = layer.pre_mlp_norm_checkpoint.checkpoint( apply_module(layer.pre_mlp_layernorm), hidden_states @@ -570,6 +653,8 @@ def forward_func( f"got {len(pre_mlp_layernorm_output)}" ) pre_mlp_layernorm_output, hidden_states = pre_mlp_layernorm_output + if not is_mhc_layer: + residual = hidden_states shared_expert_output = layer.mlp.shared_experts_compute(pre_mlp_layernorm_output) probs, routing_map = layer.mlp.route( @@ -578,9 +663,18 @@ def forward_func( local_tokens, probs = layer.mlp.preprocess( pre_mlp_layernorm_output, probs, routing_map ) + if is_mhc_layer: + return ( + residual, + local_tokens, + probs, + shared_expert_output, + mlp_h_res, + mlp_hc_h_post, + ) return hidden_states, local_tokens, probs, shared_expert_output - hidden_states, local_tokens, probs, shared_expert_output = forward_func( + forward_kwargs = dict( hidden_states=hidden_states, attention_mask=node.chunk_state.attention_mask, rotary_pos_emb=node.chunk_state.rotary_pos_emb, @@ -589,11 +683,27 @@ def forward_func( packed_seq_params=node.chunk_state.packed_seq_params, sequence_len_offset=node.chunk_state.sequence_len_offset, ) + if is_hyper_connection_layer and ( + not using_cuda_graph_replay + or CudaGraphModule.attn not in layer.config.cuda_graph_modules + ): + forward_kwargs["mhc_recompute_manager"] = mhc_recompute_manager + forward_outputs = forward_func(**forward_kwargs) + if is_mhc_layer: + hidden_states, local_tokens, probs, shared_expert_output, mlp_h_res, mlp_hc_h_post = ( + forward_outputs + ) + else: + hidden_states, local_tokens, probs, shared_expert_output = forward_outputs + mlp_h_res, mlp_hc_h_post = None, None if not isinstance(layer.mlp, MoELayer): return hidden_states # Detach here for mlp_bda residual connection node.layer_state.residual = node.detach(hidden_states) + if is_mhc_layer: + node.layer_state.mlp_h_res = node.detach(mlp_h_res) + node.layer_state.mlp_hc_h_post = node.detach(mlp_hc_h_post) if layer.mlp.use_shared_expert and not layer.mlp.shared_expert_overlap: # Detach here for shared expert connection in moe_combine node.layer_state.shared_expert_output = node.detach(shared_expert_output) @@ -649,25 +759,32 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): """ - # Triggers token combine and the remaining computation in the transformer layer. - # The `mlp_bda` computation is placed after `mlp.combine` due to data dependency. - # This ordering is also critical for pipeline performance. Starting the `mlp.combine` - # communication at first allows it to be overlapped with computation from another - # microbatch. If `mlp_bda` were to run first, it would compete for SM resources - # with another microbatch's computation and expose the communication. + Trigger token combine and the remaining layer computation. + + MHC post-processing stays in this communication-stream node so it preserves the + existing EP overlap stream topology. """ residual = node.layer_state.residual shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None) 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() + + if shared_expert_output is not None: + shared_expert_output.record_stream(torch.cuda.current_stream()) + node.layer_state.shared_expert_output = None + + if is_mhc_layer: + return submodule_mhc_post_forward(node, output) + + mlp_output_with_bias = (output, None) with layer.bias_dropout_add_exec_handler(): hidden_states = layer.mlp_bda(layer.training, layer.config.bias_dropout_fusion)( mlp_output_with_bias, residual, layer.hidden_dropout ) + # Delay the offload of the mlp norm until after the mlp_bda has been computed # because the residual is needed in the mlp_bda. mlp_norm_manager = getattr(node.layer_state, 'mlp_norm_manager', None) @@ -676,30 +793,68 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): hidden_states, forced_released_tensors=[residual] ) node.layer_state.mlp_norm_manager = None - output = make_viewless_tensor( - inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True - ) + output = finalize_decoder_layer_output(node, hidden_states) # Need to record tensors created on comp stream to comm stream node.layer_state.residual.record_stream(torch.cuda.current_stream()) - if shared_expert_output is not None: - shared_expert_output.record_stream(torch.cuda.current_stream()) # release tensor reference after use node.layer_state.residual = None - node.layer_state.shared_expert_output = None + return output + + def submodule_mhc_post_forward(node: ScheduleNode, output: torch.Tensor): + """Run MLP-side mHC post-processing after combine communication completes.""" + residual = node.layer_state.residual + manager = getattr(node, "mhc_recompute_manager", None) + is_group_end = getattr(node, "is_last_layer_in_mhc_recompute_group", False) + bda_manager = None if is_group_end else manager + hidden_states = layer._forward_mhc_mlp_post( + output, + node.layer_state.mlp_h_res, + residual, + node.layer_state.mlp_hc_h_post, + bda_manager, + ) + + mlp_norm_manager = getattr(node.layer_state, 'mlp_norm_manager', None) + if mlp_norm_manager is not None: + hidden_states = mlp_norm_manager.group_offload( + hidden_states, forced_released_tensors=[residual] + ) + node.layer_state.mlp_norm_manager = None + + output = finalize_decoder_layer_output(node, hidden_states) + + node.layer_state.residual.record_stream(torch.cuda.current_stream()) + node.layer_state.mlp_h_res.record_stream(torch.cuda.current_stream()) + node.layer_state.mlp_hc_h_post.record_stream(torch.cuda.current_stream()) + node.layer_state.residual = None + node.layer_state.mlp_h_res = None + node.layer_state.mlp_hc_h_post = None + + if manager is not None and is_group_end: + manager.discard_all_outputs() - # final layer norm from decoder - final_layernorm = node.chunk_state.model.decoder.final_layernorm - if not node.is_mtp and final_layernorm and node.is_last_layer: - output = final_layernorm(output) - output = make_viewless_tensor(inp=output, requires_grad=True, keep_graph=True) return output @copy_signature(layer._forward_mlp, handle_first_dst_param='preserve') def mlp_wrapper(node: ScheduleNode, *args, **kwargs): - """Wrapper for Dense forward.""" - return layer._forward_mlp(*args, **kwargs) + """Wrapper for dense forward with explicit mHC recompute management.""" + manager = ( + getattr(node, "mhc_recompute_manager", None) if is_hyper_connection_layer else None + ) + if manager is not None: + manager.is_last_layer_in_recompute_block = getattr( + node, "is_last_layer_in_mhc_recompute_group", False + ) + kwargs["mhc_recompute_manager"] = manager + output = layer._forward_mlp(*args, **kwargs) + # Dense layers are terminal for their own layer, so the decoder boundary (mHC + # contraction + final layer norm) must be applied here for a dense final layer. + output = finalize_decoder_layer_output(node, output) + if manager is not None and getattr(node, "is_last_layer_in_mhc_recompute_group", False): + manager.discard_all_outputs() + return output def raise_not_implemented(*args): """Raise NotImplementedError for Dense layer.""" @@ -719,12 +874,26 @@ def raise_not_implemented(*args): def build_mtp_layer_callables(layer): - """Callables for multi-token prediction layer nodes. + """Create callables for multi-token prediction layer schedule nodes. - This class contains the callable functions for different types of - multi-token prediction layer nodes (attention, MLP, etc.) - """ + The returned forward callables use the same five-slot layout as transformer layers: + 1. Attention with MTP preprocessing. + 2. MoE dispatch. + 3. MoE experts. + 4. MoE combine and MLP-side mHC post-processing. + 5. MTP post-processing. + + Args: + layer: Multi-token prediction layer whose underlying transformer layer is decomposed. + + Returns: + A tuple containing the ordered forward callables and a mapping of node names to backward + weight-gradient callables. + + Raises: + AssertionError: If the underlying transformer layer is not an MoE layer. + """ forward_funcs, backward_dw = build_transformer_layer_callables(layer.mtp_model_layer) attn_forward, dispatch_forward, mlp_forward, combine_forward, _ = forward_funcs is_moe = isinstance(layer.mtp_model_layer.mlp, MoELayer) @@ -735,7 +904,11 @@ def submodule_mtp_attn_forward(node, hidden_states): if node.is_first_layer: offset = get_mtp_layer_offset(layer.config, node.chunk_state.model.vp_stage) node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0)) - hidden_states = node.chunk_state.mtp_hidden_states[offset] + mhc_multistream = getattr(node.chunk_state, "mhc_multistream", None) + if layer.config.enable_hyper_connections and mhc_multistream is not None: + hidden_states = list(torch.chunk(mhc_multistream, 1 + offset, dim=0))[offset] + else: + hidden_states = node.chunk_state.mtp_hidden_states[offset] input_ids, position_ids, padding_mask, decoder_input, hidden_states = layer._get_embeddings( input_ids=node.chunk_state.input_ids, @@ -769,11 +942,18 @@ def submodule_mtp_attn_forward(node, hidden_states): return attn_forward(node, hidden_states) def submodule_mtp_postprocess_forward(node, hidden_states): + # Save pre-contraction multi-stream; _postprocess contracts for mtp_hidden_states. + pre_contraction_hidden_states = ( + hidden_states if layer.config.enable_hyper_connections else None + ) hidden_states = layer._postprocess(hidden_states) node.chunk_state.mtp_hidden_states.append(hidden_states) if node.is_last_layer: hidden_states = torch.cat(node.chunk_state.mtp_hidden_states, dim=0) node.chunk_state.mtp_hidden_states = None + node.chunk_state.mhc_multistream = None + elif pre_contraction_hidden_states is not None: + hidden_states = pre_contraction_hidden_states return hidden_states def rng_context_wrapper(func, *args, **kwargs): @@ -796,10 +976,17 @@ def rng_context_wrapper(func, *args, **kwargs): mtp_post_process_func = submodule_mtp_postprocess_forward forward_funcs = [attn_func, dispatch_func, mlp_func, combine_func, mtp_post_process_func] + # Under hyper-connections the MTP layer builds separate e_proj/h_proj and + # sets eh_proj to None; appending eh_proj unconditionally would place None + # in the delayed-wgrad callable list (crashing backward_dw) and leave the + # e_proj/h_proj weight gradients without a delayed-wgrad trigger. + mtp_projs = ( + [layer.e_proj, layer.h_proj] if layer.config.enable_hyper_connections else [layer.eh_proj] + ) if isinstance(backward_dw["attn"], list): - backward_dw["attn"].append(layer.eh_proj) + backward_dw["attn"].extend(mtp_projs) else: - backward_dw["attn"] = [backward_dw["attn"], layer.eh_proj] + backward_dw["attn"] = [backward_dw["attn"], *mtp_projs] return forward_funcs, backward_dw diff --git a/megatron/core/pipeline_parallel/utils.py b/megatron/core/pipeline_parallel/utils.py index 0593693501c..e5afd5ac388 100644 --- a/megatron/core/pipeline_parallel/utils.py +++ b/megatron/core/pipeline_parallel/utils.py @@ -156,6 +156,8 @@ def __init__( backward_func: Optional[Callable] = None, free_input: bool = False, name: str = "schedule_node", + forward_nvtx_name: Optional[str] = None, + backward_nvtx_name: Optional[str] = None, ): """Initialize a schedule node. @@ -173,8 +175,12 @@ def __init__( free_input (bool): Flag to indicate if the input should be freed after the forward pass. name (str): Name of the node for debugging purposes. + forward_nvtx_name (str, optional): Stable NVTX label for forward execution. + backward_nvtx_name (str, optional): Stable NVTX label for backward execution. """ self.name = name + self.forward_nvtx_name = forward_nvtx_name or f"{name} forward" + self.backward_nvtx_name = backward_nvtx_name or f"{name} backward" self.forward_func = forward_func self.backward_func = backward_func if backward_func else self.default_backward_func self.stream = stream @@ -206,7 +212,7 @@ def _forward(self, *inputs): # Lazy initialization of stream if isinstance(self.stream, Callable): self.stream = self.stream() - with self.stream_acquire_context(f"{self.name} forward"): + with self.stream_acquire_context(self.forward_nvtx_name): self.inputs = [make_viewless(e).detach() if e is not None else None for e in inputs] for i, input in enumerate(self.inputs): if input is not None: @@ -215,7 +221,9 @@ def _forward(self, *inputs): data = tuple(self.inputs) data = self.forward_func(*data) - if not isinstance(data, tuple): + if data is None: + pass + elif not isinstance(data, tuple): data = make_viewless(data) else: data = tuple([make_viewless(e) if isinstance(e, torch.Tensor) else e for e in data]) @@ -246,7 +254,7 @@ def _backward(self, *output_grad): # Lazy initialization of stream if isinstance(self.stream, Callable): self.stream = self.stream() - with self.stream_acquire_context(f"{self.name} backward"): + with self.stream_acquire_context(self.backward_nvtx_name): outputs = self.output if not isinstance(outputs, tuple): outputs = (outputs,) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 4ad95ce2253..dbecd73abd3 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -759,22 +759,29 @@ def backward(ctx, *args): class CheckpointManager: - """ - Manages multiple CheckpointWithoutOutput objects within a TransformerBlock - cross layer recomputations, enabling unified recomputation during backward pass. - This is particularly useful for scenarios where multiple checkpoint operations have - sequential dependencies (i.e., the output of one checkpoint is the input of the next). - - Usage: - ckptManager = CheckpointManager() - ckpt_function = CheckpointWithoutOutput(ckpt_manager=ckptManager) + """Manage checkpoints that are recomputed together across transformer layers. + + This manager enables unified recomputation for checkpoint operations with sequential + dependencies, such as when one checkpoint's output is the next checkpoint's input. + + Examples: + ckpt_manager = CheckpointManager() + ckpt_function = CheckpointWithoutOutput(ckpt_manager=ckpt_manager) ckpt_function.checkpoint(run_function, *args) # other checkpointed operations + + # Hook-driven path: ckpt_manager.discard_all_outputs_and_register_unified_recompute(final_output) + + # Or scheduler-driven path: + ckpt_manager.discard_all_outputs() + ckpt_manager.recompute_now() """ def __init__(self): self.checkpoints = [] + self._outputs_discarded = False + self._recomputed = False # Set by TransformerBlock before each layer forward. # When True, the layer should keep block-boundary output uncheckpointed. self.is_last_layer_in_recompute_block = False @@ -789,19 +796,42 @@ def add_checkpoint(self, ckpt): def discard_all_outputs_and_register_unified_recompute(self, hook_tensor): """Discard all checkpoint outputs to save memory and register unified recompute hook.""" - for ckpt in self.checkpoints: - for output in ckpt.outputs: - output.untyped_storage().resize_(0) + self.discard_all_outputs() # Register unified recompute hook if hook_tensor.requires_grad: hook_tensor.register_hook(self._unified_recompute_hook) - def _unified_recompute_hook(self, grad_output): + def discard_all_outputs(self) -> None: + """Discard all managed checkpoint outputs without registering a backward hook. + + This operation is idempotent; calls after the first successful discard are no-ops. + """ + if self._outputs_discarded: + return + for ckpt in self.checkpoints: + for output in ckpt.outputs: + output.untyped_storage().resize_(0) + self._outputs_discarded = True + + def recompute_now(self) -> None: + """Eagerly replay all managed checkpoints in their original forward order. + + This operation is idempotent; calls after the first successful replay are no-ops. + + Raises: + RuntimeError: If the managed outputs have not been discarded before replay. + """ + if self._recomputed: + return + if not self._outputs_discarded: + raise RuntimeError("CheckpointManager.recompute_now() requires discarded outputs.") for ckpt in self.checkpoints: - # Call _recompute for each checkpoint in forward order - # The _recompute method will restore the output tensor storage ckpt._recompute(None) + self._recomputed = True + + def _unified_recompute_hook(self, grad_output): + self.recompute_now() class CheckpointWithoutOutput(object): diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index a4404624678..79a13880804 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy from dataclasses import dataclass @@ -893,6 +893,11 @@ def __init__( submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon ) + def backward_dw(self): + """Compute the deferred weight gradients (delay_wgrad_compute) of the compressor linears.""" + self.linear_wkv.backward_dw() + self.linear_wgate.backward_dw() + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: """Apply overlapping window transform for 4x compression. @@ -1285,6 +1290,12 @@ def __init__( name=(name + ".compressor") if name is not None else None, ) + def backward_dw(self): + """Compute the deferred weight gradients (delay_wgrad_compute) of the indexer linears.""" + self.linear_wq_b.backward_dw() + self.linear_weights_proj.backward_dw() + self.compressor.backward_dw() + def forward_before_topk( self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> Union[ @@ -1549,6 +1560,17 @@ def __init__( else: self.indexer = None + def backward_dw(self): + """Compute the deferred weight gradients of the optional compressor/indexer submodules. + + The None-guards mirror __init__: compressor exists only for compress_ratio > 1, + the indexer only for compress_ratio == 4 outside csa_dense_mode. + """ + if self.compressor is not None: + self.compressor.backward_dw() + if self.indexer is not None: + self.indexer.backward_dw() + # ------------------------------------------------------------------ # Private helpers – each owns one logical slice of the forward pass. # ------------------------------------------------------------------ diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 07bb1812296..3c8f139d8a0 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -911,6 +911,10 @@ def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" self._backward_kv_proj() self._backward_q_proj() + # core_attention is always CompressedSparseAttention for the dsv4_hybrid + # variant; its compressor/indexer linears defer their wgrads under + # delay_wgrad_compute and must be flushed here as well. + self.core_attention.backward_dw() self._backward_output_proj() def _backward_kv_proj(self): diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 688219f860a..806df913204 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy import math @@ -1382,6 +1382,12 @@ def __init__( parallel_mode="duplicated", ) + def backward_dw(self): + """Compute the deferred weight gradients (delay_wgrad_compute) of the indexer linears.""" + self.linear_wq_b.backward_dw() + self.linear_wk.backward_dw() + self.linear_weights_proj.backward_dw() + def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: float): """Apply RoPE to the input tensor.""" # x_pe [seqlen, batch, *, qk_pos_emb_head_dim] @@ -1621,6 +1627,10 @@ def __init__( ) self.softmax_scale = softmax_scale + def backward_dw(self): + """Compute the deferred weight gradients (delay_wgrad_compute) of the indexer.""" + self.indexer.backward_dw() + def forward( self, query: torch.Tensor, diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index f2372f4ee42..66568163178 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import math @@ -1087,6 +1087,13 @@ def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" self._backward_kv_proj() self._backward_q_proj() + # For the 'dsa' experimental variant core_attention is DSAttention, whose + # indexer linears defer their wgrads under delay_wgrad_compute and need an + # explicit flush. For standard MLA the core is TEDotProductAttention, which + # owns no linears and defines no backward_dw — hence the guard. + core_attention_backward_dw = getattr(self.core_attention, "backward_dw", None) + if core_attention_backward_dw is not None: + core_attention_backward_dw() self._backward_output_proj() def _backward_kv_proj(self): diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index 10b1de265c8..142fba9a69f 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -442,6 +442,115 @@ def has_final_layernorm_in_this_stage(self): and self.post_layer_norm ) + def preprocess_for_layer_schedule(self, hidden_states: Union[Tensor, WrappedTensor]) -> Tensor: + """Apply TransformerBlock entry processing for normal and scheduled forward paths. + + On the first pipeline stage, this method uses the provided hidden states. On subsequent + stages, it uses the tensor supplied through :meth:`set_input_tensor`. The result is made + viewless and, on the first pipeline stage, expanded into mHC residual streams when mHC is + enabled. + + Args: + hidden_states: Hidden states for the first pipeline stage, optionally wrapped for + deferred unwrapping. + + Returns: + Viewless hidden states ready for layer execution. When mHC is enabled on the first + pipeline stage, the hidden dimension contains all residual streams. + """ + # Delete the obsolete reference to the initial input tensor if necessary. + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if not self.pre_process: + # See set_input_tensor(). + hidden_states = self.input_tensor + + # Make the input viewless. This is usually redundant — embedding outputs and + # p2p_communication.py both produce viewless tensors — but make_viewless_tensor() + # is a negligible-overhead no-op on already-viewless inputs, so it is kept here + # defensively for mbs == 1 view-tensor and other corner cases. + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + + # Expand hidden states for hyper connections at the start of the block. + # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage. + if self.config.enable_hyper_connections and self.pre_process: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.num_residual_streams + ) # [s, b, C] -> [s, b, n*C] + + return hidden_states + + def postprocess_for_layer_schedule( + self, + hidden_states: Tensor, + *, + extract_layer_indices: Optional[Set[int]] = None, + return_mhc_multistream: bool = False, + ) -> Union[Tensor, Tuple[Tensor, Optional[Tensor]]]: + """Apply TransformerBlock exit processing shared by normal and scheduled forward paths. + + This method contracts mHC residual streams on the stage that owns the final layer norm, + applies the final layer norm, and preserves a distinct output node for empty pipeline + stages. With MTP, it can also return the pre-contraction mHC streams for the MTP input. + + Args: + hidden_states: Hidden states produced by the transformer layers. + extract_layer_indices: Requested feature-extraction layer indices. Nonempty feature + extraction is not supported when mHC and MTP are both enabled. + return_mhc_multistream: Whether to return the pre-contraction mHC streams together with + the processed hidden states. + + Returns: + The processed hidden states. If ``return_mhc_multistream`` is true, returns a tuple of + the processed hidden states and the pre-contraction mHC streams. The second element is + ``None`` when no MTP mHC streams need to be preserved. + + Raises: + AssertionError: If nonempty feature extraction is requested while mHC and MTP are both + enabled. + """ + mhc_multistream = None + # Only contract if the final layer norm is in this stage. + if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): + # When MTP is enabled, save pre-contraction multi-stream for MTP input. + if self.config.mtp_num_layers is not None: + if extract_layer_indices is not None: + assert ( + len(extract_layer_indices) == 0 + ), "Feature extraction is not supported with mHC + MTP." + mhc_multistream = hidden_states + # DSv4 introduced the new output contraction for mHC. + # [s, b, n*C] -> [s, b, C] + hidden_states = learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.num_residual_streams, + self.config.layernorm_epsilon, + ) + + # Final layer norm. + if self.final_layernorm is not None: + hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) + # TENorm produces a "viewed" tensor. This will result in schedule.py's + # deallocate_output_tensor() throwing an error, so a viewless tensor is + # created to prevent this. + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True + ) + + # If this TransformerBlock is empty, input and output hidden states will be the same + # node on the computational graph and will lead to unexpected errors in pipeline + # schedules. + if not self.pre_process and len(self.layers) == 0 and not self.final_layernorm: + hidden_states = hidden_states.clone() + + if return_mhc_multistream: + return hidden_states, mhc_multistream + return hidden_states + def _setup_fused_tp_communication(self): """Setup fused TP communication for all layers. We have a fused reduce-scatter + add + layer-norm + all-gather operation. @@ -799,37 +908,7 @@ def forward( self.config, self.vp_stage, get_pg_rank(pp_group) ) - # Delete the obsolete reference to the initial input tensor if necessary - if isinstance(hidden_states, WrappedTensor): - hidden_states = hidden_states.unwrap() - - if not self.pre_process: - # See set_input_tensor() - hidden_states = self.input_tensor - - # Viewless tensor. - # - We only need to create a viewless tensor in the case of micro batch - # size (mbs) == 1, since in this case, 'hidden_states.transpose()' - # above creates a view tensor, and '.contiguous()' is a pass-through. - # For mbs >= 2, '.contiguous()' creates a new tensor, eliminating - # the need to make it viewless. - # - # However, we don't explicitly check mbs == 1 here because - # make_viewless_tensor() has negligible overhead when its input - # is already viewless. - # - # - For the 'else' case above, calling make_viewless_tensor() here is - # likely redundant, since p2p_communication.py (likely originator) - # already creates viewless tensors. That said, make_viewless_tensor() - # is called here to be future-proof and corner-case-proof. - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) - - # Expand hidden states for hyper connections at the start of the block - # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage - if self.config.enable_hyper_connections and self.pre_process: - hidden_states = HyperConnectionModule.input_expand( - hidden_states, self.num_residual_streams - ) # [s, b, C] -> [s, b, n*C] + hidden_states = self.preprocess_for_layer_schedule(hidden_states) if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() @@ -952,40 +1031,9 @@ def forward( if (l_no + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) - # Only contract if the final layer norm is in this stage - mhc_multistream = None - if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): - # When MTP is enabled, save pre-contraction multi-stream for MTP input. - if self.config.mtp_num_layers is not None: - assert ( - len(extract_layer_indices) == 0 - ), "Feature extraction is not supported with mHC + MTP." - mhc_multistream = hidden_states - # DSv4 introduced the new output contraction for mHC. - # [s, b, n*C] -> [s, b, C] - hidden_states = learned_output_contract( - hidden_states, - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - self.config.num_residual_streams, - self.config.layernorm_epsilon, - ) - - # Final layer norm. - if self.final_layernorm is not None: - hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) - # TENorm produces a "viewed" tensor. This will result in schedule.py's - # deallocate_output_tensor() throwing an error, so a viewless tensor is - # created to prevent this. - hidden_states = make_viewless_tensor( - inp=hidden_states, requires_grad=True, keep_graph=True - ) - - # If this TransformerBlock is empty, input and output hidden states will be the same node - # on the computational graph and will lead to unexpected errors in pipeline schedules. - if not self.pre_process and len(self.layers) == 0 and not self.final_layernorm: - hidden_states = hidden_states.clone() + hidden_states, mhc_multistream = self.postprocess_for_layer_schedule( + hidden_states, extract_layer_indices=extract_layer_indices, return_mhc_multistream=True + ) if len(extract_layer_indices) > 0: return hidden_states, intermediate_hidden_states diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 50eedfbcf24..43e45b733db 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1204,18 +1204,20 @@ class TransformerConfig(ModelParallelConfig): """ mhc_recompute_layer_num: Optional[int] = None - """Number of layers per MHC recompute block. - - When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer - in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final - layer in the transformer block) will: - - NOT checkpoint its final MLP BDA - - Register the unified recompute hook on its MLP BDA output - - A new CheckpointManager is created for subsequent layers - - If None, all layers in the transformer block share a single recompute block. - - Must be a positive integer when set.""" + """Number of layers in each mHC recompute group. + + Layers are grouped in their local transformer-block order. The last layer in each group leaves + its final MLP BDA output uncheckpointed and closes the current ``CheckpointManager``; a new + manager is created for the next group. + + In the standard forward path, the group-ending layer registers a unified recompute hook on its + MLP BDA output. In the fine-grained expert-parallel overlap schedule, the group outputs are + discarded explicitly and a compute-stream schedule node replays the group before its backward + computation. + + If ``None``, all local layers in the transformer block share one recompute group. The value must + be a positive integer when set. + """ #################### # miscellaneous @@ -2195,7 +2197,7 @@ def __post_init__(self): if self.fine_grained_activation_offloading and self.offload_modules: # mHC checkpoints wrap input_layernorm (inside attn_norm offload context) # and pre_mlp_layernorm (inside mlp_norm offload context). The unified - # recompute hook fires before GroupCommitFunction.backward() initializes + # recompute trigger runs before GroupCommitFunction.backward() initializes # the backward chunk, so tensor_pop hits a None chunk for these modules. # Other offload modules (qkv_linear, core_attn, attn_proj, expert_fc1, # moe_act) live inside self_attention/MLP which are NOT wrapped by mHC @@ -2205,13 +2207,26 @@ def __post_init__(self): if conflicting: raise ValueError( f"'mhc' in recompute_modules is incompatible with " - f"offload_modules {conflicting}. The mHC recompute hook fires " + f"offload_modules {conflicting}. The mHC recompute replay starts " f"before the offloading backward chunk is initialized for these " f"modules, causing tensor_pop on a None chunk. Remove " f"{conflicting} from offload_modules or remove 'mhc' from " f"recompute_modules." ) + if ( + self.overlap_moe_expert_parallel_comm + and self.recompute_granularity == "selective" + and "mhc" in self.recompute_modules + and ( + self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph + ) + ): + raise ValueError( + "mHC recompute with overlap_moe_expert_parallel_comm requires CUDA graphs " + "to be disabled because explicit group replay is eager-only." + ) + if self.enable_hyper_connections and not ( self.recompute_granularity == "selective" and "mhc" in self.recompute_modules ): diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ab98b025d62..34a456bfa69 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -2129,6 +2129,13 @@ def _forward_mlp( hidden_states, mhc_recompute_manager=mhc_recompute_manager ) nvtx_range_pop(suffix="mlp_hyper_connection") + # mHC aggregation upcasts the single-stream MLP input to fp32 for numerical + # stability of the mixing weights, and that fp32 propagates through the residual + # stream. MoE layers absorb it in their own routing dtype handling, but a dense + # MLP feeds it straight into a low-precision TE layer whose fused layernorm weight + # is in params_dtype, so restore params_dtype for dense (non-MoE) layers. + if not self.is_moe_layer and hidden_states.dtype != self.config.params_dtype: + hidden_states = hidden_states.to(self.config.params_dtype) # Optional Layer norm post the cross-attention. checkpoint_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm or ( @@ -2217,6 +2224,43 @@ def _forward_mlp( mlp_output_with_bias, mlp_h_res, residual, mlp_hc_h_post, mhc_mlp_bda_manager ) + def _forward_mhc_mlp_post_core( + self, + mlp_output_with_bias, + mlp_h_res, + residual, + mlp_hc_h_post, + mhc_mlp_bda_recompute_manager: Optional['CheckpointManager'] = None, + ): + """Run the fused mHC post-MLP H_res/H_post/BDA computation.""" + nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( + mlp_h_res, + residual, + mlp_hc_h_post, + mlp_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_mlp_bda_recompute_manager, + ) + nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + return hidden_states + + def _forward_mhc_mlp_post( + self, + mlp_output, + mlp_h_res, + residual, + mlp_hc_h_post, + mhc_mlp_bda_recompute_manager: Optional['CheckpointManager'] = None, + ): + """Run mHC post-MLP fused H_res/H_post/BDA without MLP norm offload.""" + return self._forward_mhc_mlp_post_core( + (mlp_output, None), mlp_h_res, residual, mlp_hc_h_post, mhc_mlp_bda_recompute_manager + ) + def _forward_post_mlp_with_fused_hyper_connection( self, mlp_output_with_bias, @@ -2247,19 +2291,9 @@ def _forward_post_mlp_with_fused_hyper_connection( mlp_output_with_bias[0] ) - nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") - with self.bias_dropout_add_exec_handler(): - hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( - mlp_h_res, - residual, - mlp_hc_h_post, - mlp_output_with_bias, - self.hidden_dropout, - self.training, - self.config.bias_dropout_fusion, - mhc_mlp_bda_recompute_manager, - ) - nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + hidden_states = self._forward_mhc_mlp_post_core( + mlp_output_with_bias, mlp_h_res, residual, mlp_hc_h_post, mhc_mlp_bda_recompute_manager + ) hidden_states = self.mlp_norm_manager.group_offload(hidden_states) @@ -2346,7 +2380,14 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): probs, routing_map = self.mlp.route(hidden_states) hidden_states, probs = self.mlp.preprocess(hidden_states, probs, routing_map) nvtx_range_pop(suffix="mlp") - return residual, hidden_states, probs, shared_expert_output + return ( + residual, + hidden_states, + probs, + shared_expert_output, + mlp_h_res, + mlp_hc_h_post, + ) mlp_output_with_bias = self.mlp(hidden_states) self.mlp.cudagraph_tensor_store.clear() nvtx_range_pop(suffix="mlp") @@ -2359,6 +2400,42 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): ) self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: + if self.config.overlap_moe_expert_parallel_comm: + assert ( + len(cuda_graph_output) == 1 + ), "CUDA Graph output should be the attention output." + hidden_states = cuda_graph_output.pop() + if not self.is_moe_layer: + return hidden_states, None, None, None + + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post, residual = self.mlp_hyper_connection( + hidden_states + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + + hidden_states = apply_module(self.pre_mlp_layernorm)(hidden_states) + if isinstance(hidden_states, tuple): + if len(hidden_states) != 2: + raise ValueError( + "When the output of pre_mlp_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(hidden_states)}" + ) + hidden_states, _ = hidden_states + + shared_expert_output = self.mlp.shared_experts_compute(hidden_states) + probs, routing_map = self.mlp.route(hidden_states) + hidden_states, probs = self.mlp.preprocess(hidden_states, probs, routing_map) + return ( + residual, + hidden_states, + probs, + shared_expert_output, + mlp_h_res, + mlp_hc_h_post, + ) + output = self._forward_mlp( *cuda_graph_output, input_ids=kwargs.get("input_ids", None), diff --git a/tests/unit_tests/a2a_overlap/test_mhc_schedule.py b/tests/unit_tests/a2a_overlap/test_mhc_schedule.py new file mode 100644 index 00000000000..4370166f44b --- /dev/null +++ b/tests/unit_tests/a2a_overlap/test_mhc_schedule.py @@ -0,0 +1,642 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import gc +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.models.common.model_chunk_schedule_plan import ( + TransformerLayerSchedulePlan, + TransformerModelChunkSchedulePlan, +) +from megatron.core.pipeline_parallel.utils import get_comp_stream, set_streams +from megatron.core.tensor_parallel.random import ( + CheckpointManager, + CheckpointWithoutOutput, + initialize_rng_tracker, +) +from megatron.core.transformer.module import float16_to_fp32 +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import is_te_min_version +from tests.unit_tests.a2a_overlap.utils import ( + build_gpt_model, + build_input_data, + deterministic_mode, + get_test_config, + reset_model, +) +from tests.unit_tests.test_utilities import Utils + + +@pytest.fixture +def _initialized_model_parallel(): + Utils.initialize_model_parallel() + yield + Utils.destroy_model_parallel() + + +def _make_valid_mhc_overlap_config(**overrides): + """Build the smallest config satisfying the mHC overlap prerequisites.""" + kwargs = dict( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + ffn_hidden_size=128, + enable_hyper_connections=True, + num_residual_streams=4, + recompute_granularity="selective", + recompute_modules=["mhc"], + overlap_moe_expert_parallel_comm=True, + expert_model_parallel_size=2, + num_moe_experts=4, + moe_token_dispatcher_type="alltoall", + bf16=True, + ) + kwargs.update(overrides) + return TransformerConfig(**kwargs) + + +@pytest.mark.parametrize( + "cuda_graph_kwargs", + ({"cuda_graph_impl": "local"}, {"enable_cuda_graph": True}, {"external_cuda_graph": True}), +) +def test_mhc_overlap_recompute_rejects_cuda_graphs(cuda_graph_kwargs): + with pytest.raises(ValueError, match="eager-only"): + _make_valid_mhc_overlap_config(**cuda_graph_kwargs) + + +class _RecordingNode: + def __init__(self, calls, name): + self.calls = calls + self.name = name + + def forward(self, value=None): + self.calls.append(f"{self.name}.forward") + return value + + def backward(self, value): + self.calls.append(f"{self.name}.backward") + return value + + def backward_dw(self): + self.calls.append(f"{self.name}.backward_dw") + + +class _RecordingLayer: + def __init__(self, calls, prefix): + self.calls = calls + self.config = SimpleNamespace(ep_overlap_early_attn_memory_release=False) + self.attn = _RecordingNode(calls, f"{prefix}.attn") + self.moe_dispatch = _RecordingNode(calls, f"{prefix}.moe_dispatch") + self.mlp = _RecordingNode(calls, f"{prefix}.mlp") + self.moe_combine = _RecordingNode(calls, f"{prefix}.moe_combine") + self.mhc_recompute = None + self.mtp_post_process = _RecordingNode(calls, f"{prefix}.mtp_post_process") + + def get_fp8_context(self): + return nullcontext() + + def release_state(self): + self.calls.append(f"{self.attn.name.split('.')[0]}.release_state") + + +class _RecordingChunk: + def __init__(self, calls, layers): + self.calls = calls + self.layers = layers + self.pre_process = _RecordingNode(calls, "chunk.pre_process") + self.post_process = None + self.vp_stage = 0 + + def record_current_stream(self): + self.calls.append("chunk.record_current_stream") + + def wait_current_stream(self): + self.calls.append("chunk.wait_current_stream") + + def num_layers(self): + return len(self.layers) + + def pop_layer(self): + return self.layers.pop() + + def release_state(self): + self.calls.append("chunk.release_state") + + +def _assert_called_before(calls, first, second): + assert calls.count(first) == 1 + assert calls.count(second) == 1 + assert calls.index(first) < calls.index(second), f"Expected {first} before {second}: {calls}" + + +@pytest.mark.parametrize("explicit_recompute", (False, True), ids=("eager", "recompute")) +def test_layer_schedule_orders_recompute(explicit_recompute): + calls = [] + forward_layer = _RecordingLayer(calls, "forward") + backward_layer = _RecordingLayer(calls, "backward") + if explicit_recompute: + backward_layer.mhc_recompute = _RecordingNode(calls, "backward.mhc_recompute") + + TransformerLayerSchedulePlan.run( + forward_layer, backward_layer, f_input=object(), b_grad=object(), is_last_layer_in_bwd=True + ) + + if explicit_recompute: + _assert_called_before( + calls, "backward.mhc_recompute.forward", "backward.moe_combine.backward" + ) + else: + assert "backward.mhc_recompute.forward" not in calls + + +def test_model_chunk_recompute_groups_trigger_in_reverse_order(): + calls = [] + layers = [_RecordingLayer(calls, f"layer_{index}") for index in range(5)] + # Forward groups are [0, 1], [2, 3], [4], so their explicit replay nodes + # live on the final forward layer of each group. + for group_end in (1, 3, 4): + layers[group_end].mhc_recompute = _RecordingNode(calls, f"layer_{group_end}.mhc_recompute") + chunk = _RecordingChunk(calls, layers) + + TransformerModelChunkSchedulePlan.run(None, chunk, b_grad=object()) + + recompute_calls = [call for call in calls if ".mhc_recompute.forward" in call] + assert recompute_calls == [ + "layer_4.mhc_recompute.forward", + "layer_3.mhc_recompute.forward", + "layer_1.mhc_recompute.forward", + ] + for group_end in (4, 3, 1): + _assert_called_before( + calls, + f"layer_{group_end}.mhc_recompute.forward", + f"layer_{group_end}.moe_combine.backward", + ) + + +def test_model_chunk_builds_independent_two_layer_recompute_groups(monkeypatch): + captured_extra_args = [] + + class _CapturedLayerPlan: + def __init__(self, layer, event, state, comp_stream, comm_stream, extra_args): + captured_extra_args.append(dict(extra_args)) + + monkeypatch.setattr( + "megatron.core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan", + _CapturedLayerPlan, + ) + plan = TransformerModelChunkSchedulePlan.__new__(TransformerModelChunkSchedulePlan) + plan._event = object() + plan._model_chunk_state = object() + plan._transformer_layers = [] + config = SimpleNamespace( + enable_hyper_connections=True, + recompute_granularity="selective", + recompute_modules=["mhc"], + mhc_recompute_layer_num=2, + ) + module = SimpleNamespace(config=config, layers=[object() for _ in range(5)], training=True) + + plan._build_layer_schedule_plan(module, get_comp_stream, lambda: None, module_tag="decoder") + + managers = [extra_args["mhc_recompute_manager"] for extra_args in captured_extra_args] + assert managers[0] is managers[1] + assert managers[2] is managers[3] + assert managers[0] is not managers[2] + assert managers[4] is not managers[2] + assert [ + extra_args["is_last_layer_in_mhc_recompute_group"] for extra_args in captured_extra_args + ] == [False, True, False, True, True] + assert [extra_args["mhc_recompute_group_index"] for extra_args in captured_extra_args] == [ + 0, + 0, + 1, + 1, + 2, + ] + + +def test_checkpoint_manager_explicit_recompute_is_idempotent_and_restores_gradients( + _initialized_model_parallel, +): + def run_function(value): + return torch.sin(value) * value + + initialize_rng_tracker(force_reset=True) + input_tensor = torch.randn(32, device="cuda", requires_grad=True) + reference_input = input_tensor.detach().clone().requires_grad_(True) + + reference_output = run_function(reference_input) + reference_loss = reference_output.square().sum() + reference_loss.backward() + + manager = CheckpointManager() + checkpoint = CheckpointWithoutOutput(ckpt_manager=manager) + output = checkpoint.checkpoint(run_function, input_tensor) + expected_output = output.detach().clone() + loss = output.square().sum() + + manager.discard_all_outputs() + assert output.untyped_storage().nbytes() == 0 + + manager.recompute_now() + torch.testing.assert_close(output, expected_output) + + # A second explicit trigger must be an observable no-op. + manager.recompute_now() + torch.testing.assert_close(output, expected_output) + + loss.backward() + torch.testing.assert_close(input_tensor.grad, reference_input.grad) + + +def _run_schedule_and_capture(model, data): + schedule_plan = model.build_schedule_plan(**data) + output = TransformerModelChunkSchedulePlan.run(schedule_plan, None) + output_value = output.detach().clone() + TransformerModelChunkSchedulePlan.run(None, schedule_plan, b_grad=torch.ones_like(output)) + torch.cuda.synchronize() + gradients = { + name: parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + return output_value, gradients + + +def _run_eager_and_capture(model, data): + output = float16_to_fp32(model.forward(**data)) + output_value = output.detach().clone() + output.backward(torch.ones_like(output)) + torch.cuda.synchronize() + gradients = { + name: parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + return output_value, gradients + + +def _run_interleaved_schedule_and_capture(model, batches): + first_plan = model.build_schedule_plan(**batches[0]) + first_output = TransformerModelChunkSchedulePlan.run(first_plan, None) + first_output_value = first_output.detach().clone() + second_plan = model.build_schedule_plan(**batches[1]) + second_output = TransformerModelChunkSchedulePlan.run( + second_plan, first_plan, b_grad=torch.ones_like(first_output) + ) + second_output_value = second_output.detach().clone() + TransformerModelChunkSchedulePlan.run(None, second_plan, b_grad=torch.ones_like(second_output)) + torch.cuda.synchronize() + gradients = { + name: parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + return [first_output_value, second_output_value], gradients + + +def _run_eager_batches_and_capture(model, batches): + outputs = [float16_to_fp32(model.forward(**data)) for data in batches] + output_values = [output.detach().clone() for output in outputs] + for output in outputs: + output.backward(torch.ones_like(output)) + torch.cuda.synchronize() + gradients = { + name: parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + return output_values, gradients + + +def _make_mhc_numerical_config(overlap=True, recompute=True, extra_config=None): + recompute_kwargs = ( + { + "recompute_granularity": "selective", + "recompute_modules": ["mhc"], + "mhc_recompute_layer_num": 2, + } + if recompute + else {"recompute_granularity": None, "recompute_modules": []} + ) + extra_kwargs = { + "moe_token_dispatcher_type": "alltoall", + "overlap_moe_expert_parallel_comm": overlap, + "enable_hyper_connections": True, + "mhc_sinkhorn_iterations": 5, + **recompute_kwargs, + } + if extra_config: + extra_kwargs.update(extra_config) + return get_test_config(num_layers=2, extra_kwargs=extra_kwargs) + + +def _assert_close_grads(overlap_gradients, reference_gradients, rtol=5e-3, atol=5e-3): + assert overlap_gradients.keys() == reference_gradients.keys() + for name in reference_gradients: + torch.testing.assert_close( + overlap_gradients[name], + reference_gradients[name], + rtol=rtol, + atol=atol, + msg=f"Gradient mismatch for {name}", + ) + + +class TestMhcA2AOverlapNumerics: + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=4, + ) + set_streams() + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + @pytest.mark.parametrize("recompute", (False, True), ids=("without-recompute", "recompute")) + def test_two_layer_alltoall_schedule_matches_eager_hook_path(self, recompute): + reference_config = _make_mhc_numerical_config(overlap=False, recompute=recompute) + overlap_config = _make_mhc_numerical_config(recompute=recompute) + with deterministic_mode(): + data = build_input_data(seq_len=16) + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_output, reference_gradients = _run_eager_and_capture(reference_model, data) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_output, overlap_gradients = _run_schedule_and_capture(overlap_model, data) + + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + assert overlap_gradients.keys() == reference_gradients.keys() + for name in reference_gradients: + torch.testing.assert_close( + overlap_gradients[name], + reference_gradients[name], + rtol=5e-3, + atol=5e-3, + msg=f"Gradient mismatch for {name}", + ) + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + def test_two_inflight_plans_keep_recompute_groups_independent(self): + reference_config = _make_mhc_numerical_config(overlap=False) + overlap_config = _make_mhc_numerical_config() + with deterministic_mode(): + batches = [build_input_data(seq_len=16) for _ in range(2)] + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_outputs, reference_gradients = _run_eager_batches_and_capture( + reference_model, batches + ) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_outputs, overlap_gradients = _run_interleaved_schedule_and_capture( + overlap_model, batches + ) + + for overlap_output, reference_output in zip(overlap_outputs, reference_outputs): + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + assert overlap_gradients.keys() == reference_gradients.keys() + for name in reference_gradients: + torch.testing.assert_close( + overlap_gradients[name], + reference_gradients[name], + rtol=5e-3, + atol=5e-3, + msg=f"Gradient mismatch for {name}", + ) + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + def test_dense_final_layer_schedule_matches_eager(self): + # moe_layer_freq=[1, 0] makes the last decoder layer dense, so its terminal schedule + # node is the dense MLP rather than moe_combine. The decoder boundary (mHC output + # contraction + final layer norm) must still run there; otherwise an uncontracted + # [s, b, n*h] tensor reaches postprocessing (wrong shape / values). + dense_final = {"moe_layer_freq": [1, 0]} + reference_config = _make_mhc_numerical_config( + overlap=False, recompute=False, extra_config=dense_final + ) + overlap_config = _make_mhc_numerical_config(recompute=False, extra_config=dense_final) + with deterministic_mode(): + data = build_input_data(seq_len=16) + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_output, reference_gradients = _run_eager_and_capture(reference_model, data) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_output, overlap_gradients = _run_schedule_and_capture(overlap_model, data) + + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + _assert_close_grads(overlap_gradients, reference_gradients) + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + @pytest.mark.parametrize("recompute", (False, True), ids=("without-recompute", "recompute")) + def test_mtp_schedule_matches_eager(self, recompute): + # With MTP (mtp_num_layers=1, default mtp_detach_heads=False) the decoder boundary + # produces the pre-contraction mHC multi-stream consumed by the MTP depths. It must be + # detached at its producer so MTP backward does not traverse the decoder mHC graph out + # of schedule order; this node's backward_impl reconnects the accumulated gradient. + # Parameterized over recompute to exercise mHC + MTP + EP overlap + mhc recompute. + mtp = {"mtp_num_layers": 1} + reference_config = _make_mhc_numerical_config( + overlap=False, recompute=recompute, extra_config=mtp + ) + overlap_config = _make_mhc_numerical_config(recompute=recompute, extra_config=mtp) + with deterministic_mode(): + data = build_input_data(seq_len=16) + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_output, reference_gradients = _run_eager_and_capture(reference_model, data) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_output, overlap_gradients = _run_schedule_and_capture(overlap_model, data) + + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + _assert_close_grads(overlap_gradients, reference_gradients) + + @pytest.mark.skipif( + not is_te_min_version("2.3.0"), reason="delay_wgrad_compute requires TE >= 2.3.0" + ) + def test_mtp_schedule_with_delayed_wgrad_matches_eager(self): + # Regression test for the MTP delayed-wgrad callable selection: under mHC the MTP + # layer builds separate e_proj/h_proj and sets eh_proj=None, so + # build_mtp_layer_callables must register [e_proj, h_proj] (not eh_proj) with the + # attn node's delayed-wgrad list. With the unconditional eh_proj append this test + # crashes at TransformerLayerNode.backward_dw (None.backward_dw); with a wrong-but- + # non-None list it would leave the e_proj/h_proj weight gradients permanently + # deferred, which the grad-keys equality below catches. delay_wgrad_compute + # requires overlap_moe_expert_parallel_comm, so only the overlap config carries it; + # the eager reference computes its wgrads inline. + mtp = {"mtp_num_layers": 1} + reference_config = _make_mhc_numerical_config( + overlap=False, recompute=False, extra_config=mtp + ) + overlap_config = _make_mhc_numerical_config( + recompute=False, extra_config={**mtp, "delay_wgrad_compute": True} + ) + with deterministic_mode(): + data = build_input_data(seq_len=16) + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_output, reference_gradients = _run_eager_and_capture(reference_model, data) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_output, overlap_gradients = _run_schedule_and_capture(overlap_model, data) + + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + # The MTP projections must have flushed their deferred wgrads. + assert any("e_proj.weight" in name for name in overlap_gradients), overlap_gradients.keys() + assert any("h_proj.weight" in name for name in overlap_gradients), overlap_gradients.keys() + _assert_close_grads(overlap_gradients, reference_gradients) + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + @pytest.mark.parametrize("recompute", (False, True), ids=("without-recompute", "recompute")) + def test_mtp_two_inflight_plans_match_eager(self, recompute): + # Two in-flight schedule plans verify that the per-chunk mHC multi-stream (and hence + # the reconnected MTP side gradient) stays bound to the correct microbatch. + # Parameterized over recompute so the per-group recompute managers are also verified + # to stay independent across in-flight microbatches when MTP depths are present. + mtp = {"mtp_num_layers": 1} + reference_config = _make_mhc_numerical_config( + overlap=False, recompute=recompute, extra_config=mtp + ) + overlap_config = _make_mhc_numerical_config(recompute=recompute, extra_config=mtp) + with deterministic_mode(): + batches = [build_input_data(seq_len=16) for _ in range(2)] + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_outputs, reference_gradients = _run_eager_batches_and_capture( + reference_model, batches + ) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + overlap_outputs, overlap_gradients = _run_interleaved_schedule_and_capture( + overlap_model, batches + ) + + for overlap_output, reference_output in zip(overlap_outputs, reference_outputs): + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + # The tied word-embedding gradient is the most-accumulated parameter (input + # embedding + tied lm-head + MTP re-embedding, summed over both microbatches). + # Interleaving two in-flight plans reorders that bf16 reduction relative to the + # sequential eager run, so use a looser grad tolerance here. A real microbatch- + # binding error would show up as a gross mismatch on the MTP-specific parameters + # (verified bitwise-identical), not a ~2% nudge on the shared embedding. + _assert_close_grads(overlap_gradients, reference_gradients, rtol=3e-2, atol=3e-2) + + @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0") + def test_schedule_with_full_iteration_cuda_graph_matches_eager(self): + """mHC + EP overlap + ``cuda_graph_impl='full_iteration'`` (no mhc recompute). + + The ``__post_init__`` guard added by this PR rejects CUDA graphs only when + mhc *recompute* is enabled (explicit group replay is eager-only); this test + proves the guard admits full-iteration CG + EP overlap without recompute, + and that the scheduled forward+backward step is actually capturable into a + ``torch.cuda.CUDAGraph`` and numerically faithful on replay — the core-level + equivalent of what ``FullCudaGraphWrapper`` captures in production. + + MoE runs in drop_and_pad mode: the dropless alltoall dispatcher performs a + mandatory D2H splits sync that is illegal during stream capture. The + capacity factor is sized so no token can be dropped at this scale, and both + configs share the setting so the eager reference matches numerically. + """ + # capacity_factor=4.0 -> per-expert capacity = 4.0 * 16 tokens * topk(2) / 8 + # experts = 16 = the full token count, so no token can ever be dropped and + # the padded slot count (8 * 16) matches the routing-map size exactly. + drop_and_pad = {"moe_pad_expert_input_to_capacity": True, "moe_expert_capacity_factor": 4.0} + reference_config = _make_mhc_numerical_config( + overlap=False, recompute=False, extra_config=drop_and_pad + ) + overlap_config = _make_mhc_numerical_config( + recompute=False, extra_config={**drop_and_pad, "cuda_graph_impl": "full_iteration"} + ) + # Full-iteration capture requires a graph-safe RNG tracker (production + # enforces use_te_rng_tracker with CUDA graphs): TE attention forks the + # tracker even with dropout disabled, and the default tracker's + # set_state is illegal during stream capture. + initialize_rng_tracker( + use_te_rng_tracker=True, use_cudagraphable_rng=True, force_reset=True + ) + try: + self._run_full_iteration_cuda_graph_case(reference_config, overlap_config) + finally: + initialize_rng_tracker(force_reset=True) + + def _run_full_iteration_cuda_graph_case(self, reference_config, overlap_config): + with deterministic_mode(): + data = build_input_data(seq_len=16) + reference_model = build_gpt_model(reference_config) + initial_parameters = reset_model(reference_model) + reference_output, reference_gradients = _run_eager_and_capture(reference_model, data) + del reference_model + + overlap_model = build_gpt_model(overlap_config) + reset_model(overlap_model, initial_parameters) + + def run_step(): + schedule_plan = overlap_model.build_schedule_plan(**data) + output = TransformerModelChunkSchedulePlan.run(schedule_plan, None) + TransformerModelChunkSchedulePlan.run( + None, schedule_plan, b_grad=torch.ones_like(output) + ) + return output + + # Side-stream eager warmup so lazy allocations and one-time init exist + # before capture. + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(2): + run_step() + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + # Zero gradients strictly in place: the captured accumulation kernels + # must keep writing into the same live tensors on every replay. + for parameter in overlap_model.parameters(): + if parameter.grad is not None: + parameter.grad.zero_() + + graph = torch.cuda.CUDAGraph() + torch.distributed.barrier() + torch.cuda.synchronize() + with torch.cuda.graph(graph, capture_error_mode="thread_local"): + static_output = run_step() + torch.cuda.synchronize() + torch.distributed.barrier() + + # Capture only records the kernels; the first replay produces the + # real values in the static output/grad buffers. + graph.replay() + torch.cuda.synchronize() + + overlap_output = static_output.detach().clone() + overlap_gradients = { + name: parameter.grad.detach().clone() + for name, parameter in overlap_model.named_parameters() + if parameter.grad is not None + } + del graph + gc.collect() + + torch.testing.assert_close(overlap_output, reference_output, rtol=5e-3, atol=5e-3) + _assert_close_grads(overlap_gradients, reference_gradients) diff --git a/tests/unit_tests/a2a_overlap/utils.py b/tests/unit_tests/a2a_overlap/utils.py index c4d6a2844e1..9b664d4c122 100644 --- a/tests/unit_tests/a2a_overlap/utils.py +++ b/tests/unit_tests/a2a_overlap/utils.py @@ -291,11 +291,23 @@ def apply_flex_backend_kwargs(extra_kwargs, dispatcher_type, flex_backend): def build_gpt_model(config, vocab_size=512, max_seq_len=300): - """Build and return a GPTModel on CUDA from the given config.""" + """Build and return a GPTModel on CUDA from the given config. + + When ``config.mtp_num_layers`` is set, the MTP block spec is built and passed so the + model has a real multi-token-prediction block (needed to exercise the mHC multi-stream + path shared between the decoder boundary and MTP depths). + """ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.core.models.gpt.gpt_model import GPTModel layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True) + mtp_block_spec = None + if config.mtp_num_layers: + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec + + mtp_block_spec = get_gpt_mtp_block_spec( + config=config, spec=layer_spec, use_transformer_engine=True + ) model = GPTModel( config=config, transformer_layer_spec=layer_spec, @@ -303,6 +315,7 @@ def build_gpt_model(config, vocab_size=512, max_seq_len=300): pre_process=True, post_process=True, max_sequence_length=max_seq_len, + mtp_block_spec=mtp_block_spec, ) model.cuda() return model diff --git a/tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py b/tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py index eda8ffe7df4..4d31ba8a0fd 100644 --- a/tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py +++ b/tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py @@ -24,9 +24,10 @@ import torch from megatron.core import parallel_state +from megatron.core.models.gpt.fine_grained_callables import PostProcessNode, PreProcessNode from megatron.core.pipeline_parallel.schedules import get_tensor_shapes from megatron.core.transformer.hyper_connection import HyperConnectionModule -from megatron.core.transformer.transformer_block import get_num_layers_to_build +from megatron.core.transformer.transformer_block import TransformerBlock, get_num_layers_to_build from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -114,6 +115,35 @@ def _make_config( return TransformerConfig(**kwargs) +def _make_boundary_block( + config, + *, + pre_process=True, + final_layernorm=None, + has_final_layernorm=False, + num_layers=1, + input_tensor=None, +): + """Create a lightweight TransformerBlock instance for boundary helper tests.""" + block = TransformerBlock.__new__(TransformerBlock) + torch.nn.Module.__init__(block) + block.config = config + block.pre_process = pre_process + block.input_tensor = input_tensor + block.num_residual_streams = config.num_residual_streams + block.final_layernorm = final_layernorm + block.layers = [object()] * num_layers + block.has_final_layernorm_in_this_stage = lambda: has_final_layernorm + if config.enable_hyper_connections: + hidden_size = config.hidden_size + n_streams = config.num_residual_streams + device = input_tensor.device if input_tensor is not None else torch.cuda.current_device() + block.hc_head_fn = torch.randn(n_streams, n_streams * hidden_size, device=device) + block.hc_head_base = torch.zeros(n_streams, device=device) + block.hc_head_scale = torch.ones(1, device=device) + return block + + # =========================================================================== # 1. get_tensor_shapes — shape correctness with mHC # =========================================================================== @@ -403,6 +433,128 @@ def test_expand_then_contract_preserves_shape(self): # expand copies all streams → mean of identical streams = original torch.testing.assert_close(contracted, x) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_schedule_preprocess_helper_expands_first_pp_stage(self): + n = 4 + s, b, C = 8, 2, 64 + cfg = _make_config( + hidden_size=C, pp_size=2, enable_hyper_connections=True, num_residual_streams=n + ) + x = torch.randn(s, b, C, device='cuda') + block = _make_boundary_block(cfg, pre_process=True, input_tensor=x) + + expanded = block.preprocess_for_layer_schedule(x) + + assert expanded.shape == (s, b, n * C) + for stream_idx in range(n): + torch.testing.assert_close(expanded[:, :, stream_idx * C : (stream_idx + 1) * C], x) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_schedule_preprocess_helper_does_not_reexpand_non_first_pp_stage(self): + n = 4 + s, b, C = 8, 2, 64 + cfg = _make_config( + hidden_size=C, pp_size=2, enable_hyper_connections=True, num_residual_streams=n + ) + received = torch.randn(s, b, n * C, device='cuda') + block = _make_boundary_block( + cfg, pre_process=False, input_tensor=received, has_final_layernorm=False + ) + + out = block.preprocess_for_layer_schedule(torch.empty(s, b, C, device='cuda')) + + assert out.shape == (s, b, n * C) + torch.testing.assert_close(out, received) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_schedule_postprocess_helper_contracts_before_final_layernorm(self): + n = 4 + s, b, C = 8, 2, 64 + cfg = _make_config( + hidden_size=C, pp_size=2, enable_hyper_connections=True, num_residual_streams=n + ) + cfg.mtp_num_layers = 1 + multistream = torch.randn(s, b, n * C, device='cuda') + block = _make_boundary_block( + cfg, + pre_process=False, + final_layernorm=torch.nn.Identity(), + has_final_layernorm=True, + input_tensor=multistream, + ) + + contracted, saved_multistream = block.postprocess_for_layer_schedule( + multistream, return_mhc_multistream=True + ) + + assert contracted.shape == (s, b, C) + assert saved_multistream is multistream + + def test_preprocess_node_uses_block_boundary_helper(self): + decoder_input = torch.randn(8, 2, 64) + expanded_input = torch.randn(8, 2, 256) + decoder = SimpleNamespace( + input_tensor=None, preprocess_for_layer_schedule=MagicMock(return_value=expanded_input) + ) + gpt_model = SimpleNamespace( + pre_process=True, + decoder=decoder, + _preprocess=MagicMock( + return_value=(decoder_input, "rotary", "cos", "sin", "seq_offset", "pad_mask") + ), + ) + chunk_state = SimpleNamespace( + input_ids=torch.ones(2, 8, dtype=torch.long), + position_ids=torch.arange(8).repeat(2, 1), + decoder_input=None, + packed_seq_params=None, + padding_mask=None, + ) + node = PreProcessNode.__new__(PreProcessNode) + node.gpt_model = gpt_model + node.chunk_state = chunk_state + + out = node.forward_impl() + + assert out is expanded_input + assert chunk_state.decoder_input is decoder_input + assert decoder.preprocess_for_layer_schedule.call_args.args[0] is decoder_input + + def test_empty_decoder_postprocess_node_uses_block_boundary_helper(self): + hidden_states = torch.randn(8, 2, 256) + contracted = torch.randn(8, 2, 64) + loss = torch.randn(8, 2) + decoder = SimpleNamespace( + layers=[], postprocess_for_layer_schedule=MagicMock(return_value=contracted) + ) + gpt_model = SimpleNamespace(decoder=decoder, _postprocess=MagicMock(return_value=loss)) + chunk_state = SimpleNamespace( + input_ids=torch.ones(2, 8, dtype=torch.long), + position_ids=torch.arange(8).repeat(2, 1), + labels=torch.ones(2, 8, dtype=torch.long), + decoder_input=torch.randn(8, 2, 64), + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + loss_mask=None, + attention_mask=None, + packed_seq_params=None, + sequence_len_offset=None, + runtime_gather_output=None, + extra_block_kwargs=None, + output_processor=None, + output_processor_context=None, + ) + node = PostProcessNode.__new__(PostProcessNode) + node.gpt_model = gpt_model + node.chunk_state = chunk_state + + out = node.forward_impl(hidden_states) + + assert out is loss + decoder.postprocess_for_layer_schedule.assert_called_once_with(hidden_states) + assert gpt_model._postprocess.call_args.kwargs["hidden_states"] is contracted + # =========================================================================== # 3b. Zero-layer VP stage edge cases with mHC @@ -754,6 +906,10 @@ def _run_forward( use_cpu_initialization=True, pipeline_dtype=torch.bfloat16, bf16=True, + # params_dtype must match the actual .bfloat16() model params below; + # leaving it at the fp32 default makes the dense-mHC params_dtype + # cast in _forward_mlp a no-op and TE rejects the fp32 activations. + params_dtype=torch.bfloat16, pipeline_model_parallel_size=pp_size, virtual_pipeline_model_parallel_size=vp_size, enable_hyper_connections=enable_mhc, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_delay_wgrad.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_delay_wgrad.py new file mode 100644 index 00000000000..fbb3d091cf4 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_delay_wgrad.py @@ -0,0 +1,449 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Deferred weight-gradient (``delay_wgrad_compute``) flush tests for the CSA and DSA +experimental attention variants. + +Under ``delay_wgrad_compute=True`` every TE linear built from the shared config defers +its weight gradient until an explicit ``backward_dw()`` call. These tests verify that + +* ``DSv4HybridSelfAttention.backward_dw()`` traverses ``core_attention`` + (``CompressedSparseAttention`` -> ``Compressor`` / ``CSAIndexer``) so the six CSA + compressor/indexer linears of a ratio-4 layer are flushed, and +* ``MLASelfAttention.backward_dw()`` traverses ``core_attention`` (``DSAttention`` -> + ``DSAIndexer``) so the three DSA indexer linears are flushed, + +i.e. after ``loss.backward()`` the deferred weights still have ``.grad is None`` and only +after ``attn.backward_dw()`` do all of them (nested ones included) receive gradients. +""" + +import operator +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.utils import is_te_min_version +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + +_SEED = 42 + + +def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Patch hadamard_transform in dsa/csa modules if the library is not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Config / build helpers +# --------------------------------------------------------------------------- + + +def _enable_delay_wgrad_compute(config): + """Turn on delayed wgrad on an already-constructed config. + + Set AFTER construction on purpose: ``TransformerConfig.__post_init__`` couples + ``delay_wgrad_compute`` to ``overlap_moe_expert_parallel_comm`` (EP pipelining), + which is irrelevant for a single-GPU attention-module test, so we bypass that + validation by mutating the finished dataclass. This still happens BEFORE + ``build_module``: TELinear reads the flag at ``__init__`` + (megatron/core/extensions/transformer_engine.py:797-801, TE >= 2.3.0 only) and + ``TELinear.backward_dw`` gates on it (transformer_engine.py:986-989). + """ + config.delay_wgrad_compute = True + return config + + +def _make_csa_config(delay_wgrad=True): + """MLATransformerConfig for the dsv4_hybrid (CSA) delayed-wgrad test. + + Mirrors ``_make_config`` in test_dsv4_hybrid_attention.py. + """ + config = MLATransformerConfig( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=1, + sequence_parallel=False, + q_lora_rank=64, + kv_lora_rank=32, # v_head_dim - qk_pos_emb_head_dim + qk_head_dim=32, # v_head_dim - qk_pos_emb_head_dim + qk_pos_emb_head_dim=32, + v_head_dim=64, + o_groups=8, + o_lora_rank=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=[0, 4, 128, 4], + csa_window_size=8, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + # Indexer top-k is non-differentiable: the indexer linears ONLY receive + # gradients through the indexer KL loss. coeff=0 would make this test vacuous. + dsa_indexer_loss_coeff=1.0, + # Deferred wgrads must land in param.grad, not main_grad. + gradient_accumulation_fusion=False, + ) + return _enable_delay_wgrad_compute(config) if delay_wgrad else config + + +def _make_dsa_config(delay_wgrad=True): + """MLATransformerConfig for the dsa delayed-wgrad test. + + Mirrors the configs in test_attention_variant_dsa.py (TestDSAttention / + TestDSAModuleSpecDispatch) plus ``experimental_attention_variant='dsa'``. + """ + config = MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + # MLA specific configs + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=64, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + experimental_attention_variant='dsa', + # Sparse attention specific configs + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + # Indexer top-k is non-differentiable: the indexer linears ONLY receive + # gradients through the indexer KL loss. coeff=0 would make this test vacuous. + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=False, + # Deferred wgrads must land in param.grad, not main_grad. + gradient_accumulation_fusion=False, + ) + return _enable_delay_wgrad_compute(config) if delay_wgrad else config + + +def _build_csa_attention(config, layer_number, pg_collection): + """Instantiate a DSv4HybridSelfAttention from config (mirrors + test_dsv4_hybrid_attention.py::_build_attention).""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + from megatron.core.transformer.spec_utils import build_module + + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + return build_module(spec, config=config, layer_number=layer_number, pg_collection=pg_collection) + + +def _build_dsa_attention(config, layer_number, pg_collection): + """Instantiate an MLASelfAttention with a DSAttention core from config.""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsa_module_spec_for_backend, + ) + from megatron.core.transformer.spec_utils import build_module + + spec = get_dsa_module_spec_for_backend(config=config, backend=TESpecProvider()) + return build_module(spec, config=config, layer_number=layer_number, pg_collection=pg_collection) + + +# --------------------------------------------------------------------------- +# Assertion helpers +# --------------------------------------------------------------------------- + + +def _assert_flushed_grads_match_eager( + delayed_module, + delayed_named_linears, + make_config, + build_attention, + layer_number, + hidden, + pg_collection, +): + """Delay-vs-eager gradient equality: the flush must not change the math. + + Rebuild the same attention with delay_wgrad_compute off, copy the delayed + module's exact weights (state_dict, so construction-time RNG consumption is + irrelevant), replay a clone of the identical input, and require value-level + agreement of every flushed weight gradient with the inline (eager) wgrads. + This pins that backward_dw() only changes WHEN the wgrad GEMM runs, never + the gradient definition. + """ + delayed_grads = { + name: module.weight.grad.detach().clone() for name, module in delayed_named_linears + } + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + eager_attn = build_attention( + make_config(delay_wgrad=False), layer_number=layer_number, pg_collection=pg_collection + ).cuda() + eager_attn.train() + eager_attn.load_state_dict(delayed_module.state_dict()) + + eager_hidden = hidden.detach().clone().requires_grad_(True) + output, _ = eager_attn(hidden_states=eager_hidden, attention_mask=None) + output.sum().backward() + + for name, _ in delayed_named_linears: + eager_linear = operator.attrgetter(name)(eager_attn) + assert eager_linear.weight.grad is not None, f"eager reference has no grad for {name}" + torch.testing.assert_close( + delayed_grads[name], + eager_linear.weight.grad, + msg=f"delayed-flushed wgrad differs from eager inline wgrad for {name}", + ) + + +def _assert_all_grads_none(named_linears, stage): + for name, linear in named_linears: + assert linear.weight.grad is None, ( + f"{stage}: {name}.weight.grad should still be deferred (None) under " + f"delay_wgrad_compute, got a tensor" + ) + + +def _assert_all_grads_present(named_linears, stage): + for name, linear in named_linears: + assert ( + linear.weight.grad is not None + ), f"{stage}: {name}.weight.grad is None — backward_dw() did not flush this linear" + assert torch.isfinite(linear.weight.grad).all(), f"{stage}: non-finite grad on {name}" + + +def _assert_any_grad_nonzero(named_linears, stage): + assert any( + linear.weight.grad.abs().sum().item() > 0 for _, linear in named_linears + ), f"{stage}: all flushed indexer grads are exactly zero — indexer loss did not contribute" + + +# =========================================================================== +# TEST 1: dsv4_hybrid / CSA — six nested compressor/indexer linears +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +@pytest.mark.skipif( + not is_te_min_version("2.3.0"), + reason="delay_wgrad_compute requires TE >= 2.3.0 (older TE ignores the flag)", +) +class TestDSv4HybridCSADelayedWgradFlush: + """DSv4HybridSelfAttention.backward_dw() must flush the CSA nested linears.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_csa_deferred_wgrads_flushed_through_core_attention(self): + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_csa_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + # layer_number=2 -> csa_compress_ratios[1] == 4: compressor AND indexer built. + attn = _build_csa_attention(config, layer_number=2, pg_collection=pg).cuda() + attn.train() + + core = attn.core_attention + assert core.compressor is not None, "ratio-4 layer must build the compressor" + assert core.indexer is not None, "ratio-4 layer must build the indexer" + + # The six CSA linears the HEAD commit flushes through core_attention. + csa_nested_linears = [ + ("core_attention.compressor.linear_wkv", core.compressor.linear_wkv), + ("core_attention.compressor.linear_wgate", core.compressor.linear_wgate), + ("core_attention.indexer.linear_wq_b", core.indexer.linear_wq_b), + ("core_attention.indexer.linear_weights_proj", core.indexer.linear_weights_proj), + ("core_attention.indexer.compressor.linear_wkv", core.indexer.compressor.linear_wkv), + ( + "core_attention.indexer.compressor.linear_wgate", + core.indexer.compressor.linear_wgate, + ), + ] + # The attention-level TE linears also defer (flushed by backward_dw even + # before the core_attention traversal was added). + attention_level_linears = [ + ("linear_q_down_proj", attn.linear_q_down_proj), + ("linear_q_up_proj", attn.linear_q_up_proj), + ("linear_kv_proj", attn.linear_kv_proj), + ("linear_proj", attn.linear_proj), + ] + + hidden = ( + torch.randn(seq_len, batch_size, config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, config.hidden_size) + # The indexer KL loss is attached to `output` via DSAIndexerLossAutoScaler, + # so backward through the output alone reaches the indexer subgraph. + output.sum().backward() + + # dgrads flow normally; only the TE-linear wgrads are deferred. + assert hidden.grad is not None, "no dgrad on hidden_states" + # linear_o_group_proj is a raw nn.Parameter (einsum), not a TE linear: it + # receives its gradient through plain autograd, proving backward ran. + assert attn.linear_o_group_proj.grad is not None, "no grad on linear_o_group_proj" + + _assert_all_grads_none(csa_nested_linears + attention_level_linears, "pre-flush") + + attn.backward_dw() + + _assert_all_grads_present(csa_nested_linears + attention_level_linears, "post-flush") + # The four indexer-side linears receive gradients exclusively through the + # indexer loss (x/qr are detached); make sure that path is non-trivial. + _assert_any_grad_nonzero(csa_nested_linears[2:], "post-flush") + + # Same math, different timing: flushed wgrads must equal eager wgrads. + _assert_flushed_grads_match_eager( + attn, + csa_nested_linears + attention_level_linears, + _make_csa_config, + _build_csa_attention, + 2, + hidden, + pg, + ) + + +# =========================================================================== +# TEST 2: dsa — three DSA indexer linears under MLASelfAttention +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +@pytest.mark.skipif( + not is_te_min_version("2.3.0"), + reason="delay_wgrad_compute requires TE >= 2.3.0 (older TE ignores the flag)", +) +class TestDSADelayedWgradFlush: + """MLASelfAttention.backward_dw() must flush the DSA indexer linears.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_dsa_deferred_wgrads_flushed_through_core_attention(self): + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_dsa_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_dsa_attention(config, layer_number=1, pg_collection=pg).cuda() + attn.train() + + indexer = attn.core_attention.indexer + # The three DSA indexer linears the HEAD commit flushes through core_attention. + dsa_indexer_linears = [ + ("core_attention.indexer.linear_wq_b", indexer.linear_wq_b), + ("core_attention.indexer.linear_wk", indexer.linear_wk), + ("core_attention.indexer.linear_weights_proj", indexer.linear_weights_proj), + ] + # The MLA-level TE linears also defer (already flushed before the fix). + attention_level_linears = [ + ("linear_q_down_proj", attn.linear_q_down_proj), + ("linear_q_up_proj", attn.linear_q_up_proj), + ("linear_kv_down_proj", attn.linear_kv_down_proj), + ("linear_kv_up_proj", attn.linear_kv_up_proj), + ("linear_proj", attn.linear_proj), + ] + + hidden = ( + torch.randn(seq_len, batch_size, config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + # attn_mask_type=causal comes from the spec params; DSAttention then builds + # its own causal float mask, so attention_mask=None is valid here. + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, config.hidden_size) + # The indexer KL loss is attached to `output` via DSAIndexerLossAutoScaler, + # so backward through the output alone reaches the (detached-x) indexer. + output.sum().backward() + + assert hidden.grad is not None, "no dgrad on hidden_states" + # k_norm is a TE norm (not a TE linear): its weight grad arrives through + # plain autograd during the indexer-loss backward, proving that backward + # reached the indexer while the linears' wgrads stayed deferred. + assert ( + indexer.k_norm.weight.grad is not None + ), "indexer loss backward did not reach the indexer (k_norm has no grad)" + + _assert_all_grads_none(dsa_indexer_linears + attention_level_linears, "pre-flush") + + attn.backward_dw() + + _assert_all_grads_present(dsa_indexer_linears + attention_level_linears, "post-flush") + # The indexer linears receive gradients exclusively through the indexer + # loss (x/qr are detached in DSAttention.forward); ensure it contributed. + _assert_any_grad_nonzero(dsa_indexer_linears, "post-flush") + + # Same math, different timing: flushed wgrads must equal eager wgrads. + _assert_flushed_grads_match_eager( + attn, + dsa_indexer_linears + attention_level_linears, + _make_dsa_config, + _build_dsa_attention, + 1, + hidden, + pg, + ) diff --git a/tests/unit_tests/transformer/test_submodule_callables.py b/tests/unit_tests/transformer/test_submodule_callables.py index 7b41b3ca197..ecbefe50d63 100644 --- a/tests/unit_tests/transformer/test_submodule_callables.py +++ b/tests/unit_tests/transformer/test_submodule_callables.py @@ -75,6 +75,7 @@ def run_model_submodules_with_capture(model, input_tensors, microbatches): # build mock func/state node = DummyNode() node.is_mtp = False + node.is_last_layer = False node.chunk_state.model = dummy_model # attn fwd