fix(cuda_graphs): handle TE 2.15 removal of FP8GlobalStateManager.set_skip_fp8_weight_update_tensor

megatron/core/transformer/cuda_graphs.py

@@ -68,6 +68,31 @@
 _IS_GRAPH_WARMUP = False
 logger = logging.getLogger(__name__)
 
+
+def _set_skip_fp8_weight_update_tensor(skip: bool) -> None:
+    """Toggle TE's FP8 "skip weight refresh" flag between microbatches.
+
+    TE 2.15 (PR #2759) moved this flag from a classmethod on FP8GlobalStateManager
+    to a dataclass field on FP8GlobalStateManager.quantization_state, with no
+    shim. This helper handles both layouts.
+    """
+    if not HAVE_TE_GRAPHS:
+        return
+
+    # TE <= 2.14: classmethod setter still exists.
+    if hasattr(FP8GlobalStateManager, "set_skip_fp8_weight_update_tensor"):
+        FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(skip)
+        return
+
+    # TE >= 2.15: write the dataclass field directly. Allocate the persistent
+    # 1-element CUDA scalar once so its data pointer stays valid across
+    # cudagraph replays, then fill the value in place.
+    qstate = FP8GlobalStateManager.quantization_state
+    if qstate.skip_fp8_weight_update_tensor is None:
+        qstate.skip_fp8_weight_update_tensor = torch.empty(1, dtype=torch.float32, device="cuda")
+    qstate.skip_fp8_weight_update_tensor.fill_(skip)
+
+
 # Freeze GC during capture.
 # TODO (@lmcafee): remove all freeze-GC code once most users are on PyTorch 2.9+.
 FREEZE_GC = os.getenv("CUDA_GRAPH_CAPTURE_FREEZE_GC") != "0"
@@ -608,7 +633,7 @@ def forward(ctx, runner, is_first_microbatch, *inputs):
             # Note that FP8GlobalStateManager.is_first_fp8_module() is inacccurate as each
             # layer may be in its own fp8 context, when the fp8 recipe != delayed_scaling
             if runner.is_first_layer and (runner.fp8_param_cache_updated != is_first_microbatch):
-                FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(not is_first_microbatch)
+                _set_skip_fp8_weight_update_tensor(not is_first_microbatch)
                 runner.fp8_param_cache_updated = is_first_microbatch
 
         runner.fwd_graph.replay()
@@ -738,13 +763,13 @@ def __init__(
 
             if self.fp8_enabled:
                 self.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
-                FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
+                _set_skip_fp8_weight_update_tensor(False)
 
             if self.fp4_enabled:
                 from megatron.core.fp4_utils import get_fp4_recipe  # to avoid circular import
 
                 self.fp4_recipe = get_fp4_recipe(self.base_module.config)
-                FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
+                _set_skip_fp8_weight_update_tensor(False)
 
     def __str__(self):
         return "%s; hid %s" % (

tests/test_utils/recipes/h100/gpt-static-inference.yaml

@@ -68,7 +68,7 @@ products:
   - test_case: [gpt_static_inference_tp1_pp1_583m_fp8_cudagraphs]
     products:
       - environment: [dev]
-        scope: [mr-broken, mr-github-broken]
+        scope: [mr, mr-github]
         platforms: [dgx_h100]
   - test_case: [gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch]
     products:

tests/unit_tests/transformer/test_cuda_graphs.py

@@ -1527,6 +1527,49 @@ def test_num_warmup_steps_override(self):
         ), f"Expected 0 warmup steps (manager override), got {runner.num_warmup_steps}"
 
 
+class TestSkipFp8WeightUpdateTensor:
+    """Regression test for the TE 2.15 ``set_skip_fp8_weight_update_tensor`` removal."""
+
+    @staticmethod
+    def _read_skip_tensor():
+        from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+
+        getter = getattr(FP8GlobalStateManager, "get_skip_fp8_weight_update_tensor", None)
+        if getter is not None:
+            return getter()
+        return FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
+
+    @staticmethod
+    def _reset_skip_tensor():
+        from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+
+        if "skip_fp8_weight_update_tensor" in vars(FP8GlobalStateManager):
+            FP8GlobalStateManager.skip_fp8_weight_update_tensor = None
+        qstate = getattr(FP8GlobalStateManager, "quantization_state", None)
+        if qstate is not None and hasattr(qstate, "skip_fp8_weight_update_tensor"):
+            qstate.skip_fp8_weight_update_tensor = None
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+    def test_sets_value_in_place(self):
+        """Helper writes the right value and reuses the same storage across calls."""
+        from megatron.core.transformer.cuda_graphs import _set_skip_fp8_weight_update_tensor
+
+        self._reset_skip_tensor()
+        try:
+            _set_skip_fp8_weight_update_tensor(True)
+            t = self._read_skip_tensor()
+            assert t.shape == (1,) and t.dtype == torch.float32 and t.is_cuda
+            assert t.item() == 1.0
+
+            # data_ptr must stay stable so captured cudagraphs read the same address.
+            ptr = t.data_ptr()
+            _set_skip_fp8_weight_update_tensor(False)
+            assert self._read_skip_tensor().data_ptr() == ptr
+            assert self._read_skip_tensor().item() == 0.0
+        finally:
+            self._reset_skip_tensor()
+
+
 if __name__ == "__main__":
 
     test = TestParallelTransformerBlockCudagraphs()