Enable local whole-layer CUDA Graph with fine grained activation offload

megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -270,7 +270,14 @@ def forward(
             cp_rank,
             cp_size,
         )
-        ctx.save_for_backward(cos, sin)
+        # Clone cos/sin before saving: the rotary_pos_emb table is shared across
+        # the Q and KV RoPE autograd functions (same tensor passed to both), and
+        # an in-place op elsewhere in the captured layer bumps the table's
+        # autograd version after this save, making backward raise "variable
+        # needed for gradient computation has been modified by an inplace
+        # operation". Saving a private clone detaches our saved copy from any
+        # later in-place mutation of the shared table.
+        ctx.save_for_backward(cos.clone(), sin.clone())
         ctx.qk_head_dim = qk_head_dim
         ctx.emb_dim = emb_dim
         ctx.cu_seqlens_q = cu_seqlens_q
@@ -658,7 +665,11 @@ def forward(
             cp_rank,
             cp_size,
         )
-        ctx.save_for_backward(cos, sin)
+        # Clone cos/sin before saving -- see the same note in
+        # ApplyMLARotaryEmbQ.forward. The KV path and the Q path share one
+        # rotary_pos_emb table; without a clone a later in-place op on the
+        # shared table breaks this backward too.
+        ctx.save_for_backward(cos.clone(), sin.clone())
         ctx.rotary_interleaved = rotary_interleaved
         ctx.emb_dim = emb_dim
         ctx.k_dim = k_dim

megatron/core/models/gpt/gpt_layer_specs.py

@@ -38,6 +38,7 @@
     MlpBuilder,
     TransformerLayer,
     TransformerLayerSubmodules,
+    MoETransformerLayer,
     get_transformer_layer_offset,
 )
 from megatron.core.typed_torch import copy_signature, not_none
@@ -640,6 +641,10 @@ def get_gpt_decoder_layer_specs(
             kitchen_attention_backend=config.kitchen_attention_backend,
         )
 
+    # MoE layers need the specialized transformer wrapper so local partial
+    # CUDA Graphs can split router/postprocess from eager dispatch.
+    moe_layer_spec.module = MoETransformerLayer
+
     # Parse config.moe_layer_freq to determine the pattern of expert/dense layers.
     # 0 stands for dense layers, 1 stands for expert layers.
     # For integer N: Creates a pattern with one expert layer every N layers.
@@ -767,7 +772,7 @@ def get_gpt_mtp_block_spec_for_backend(
     if isinstance(spec, TransformerBlockSubmodules):
         # get the spec for the last layer of decoder block
         transformer_layer_spec = spec.layer_specs[-1]
-    elif isinstance(spec, ModuleSpec) and spec.module == TransformerLayer:
+    elif isinstance(spec, ModuleSpec) and issubclass(spec.module, TransformerLayer):
         transformer_layer_spec = spec
     else:
         raise ValueError(f"Invalid spec: {spec}")

megatron/core/pipeline_parallel/schedules.py

@@ -328,17 +328,24 @@ def forward_step_calc_loss(
             outputs = loss_func(output_tensor)
             if len(outputs) == 3:
                 output_tensor, num_tokens, loss_reduced = outputs
-                if not config.calculate_per_token_loss:
+                if not config.calculate_per_token_loss and not getattr(
+                    output_tensor, '_is_loss_head_final', False
+                ):
                     # Protect against division by zero when all tokens are masked
                     #   in a microbatch.
-                    output_tensor /= torch.clamp(num_tokens, min=1)
-                    output_tensor /= num_microbatches
+                    # NOTE: non-inplace: with the graphed loss head the loss
+                    # scalar is the output of _CudagraphRecordNode (a cudagraph
+                    # buffer view); inplace div on it is rejected by autograd.
+                    # When _is_loss_head_final is set the normalization was
+                    # already computed inside the loss-head cuda graph.
+                    output_tensor = output_tensor / torch.clamp(num_tokens, min=1)
+                    output_tensor = output_tensor / num_microbatches
             else:
                 # preserve legacy loss averaging behavior (ie, over the number of microbatches)
                 assert len(outputs) == 2
                 output_tensor, loss_reduced = outputs
-                output_tensor *= cp_group_size
-                output_tensor /= num_microbatches
+                output_tensor = output_tensor * cp_group_size
+                output_tensor = output_tensor / num_microbatches
             forward_data_store.append(loss_reduced)
         else:
             data = loss_func(output_tensor, non_loss_data=True)

megatron/core/tensor_parallel/layers.py

@@ -349,7 +349,7 @@ def forward(self, input_):
             weight = GTPEmbeddingWeight.apply(self.weight)
 
         # Get the embeddings.
-        if self.deterministic_mode:
+        if self.deterministic_mode or torch.cuda.is_available():
             output_parallel = weight[masked_input]
         else:
             # F.embedding currently has a non-deterministic backward function