diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 4b1dc3260ab..baa295b8ae8 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/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" % ( diff --git a/tests/test_utils/recipes/h100/gpt-static-inference.yaml b/tests/test_utils/recipes/h100/gpt-static-inference.yaml index de9b0235203..87046588b2b 100644 --- a/tests/test_utils/recipes/h100/gpt-static-inference.yaml +++ b/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: diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 5d2434165f5..b9e9c1ae642 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/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()