From 184a979b9d539b4754afa2b6aeb221ce081dd5ae Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Sun, 8 Nov 2026 08:26:35 -0600 Subject: [PATCH 1/3] Avoid redundant CUDA graph RNG registration Probe the installed PyTorch CUDAGraph registration API once and skip explicit per-graph registration only when it reports that registration is deprecated and is now a no-op. Preserve the existing path on builds that still require registration. Signed-off-by: RunguoLi --- tests/pytorch/test_cuda_graphs.py | 34 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 16 +++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 46aab13075c..b263d1f0b18 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -6,6 +6,7 @@ import pytest import copy import gc +import warnings import weakref import torch @@ -28,6 +29,7 @@ ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager import transformer_engine.pytorch.ops as te_ops +import transformer_engine.pytorch.graph as graph_module from transformer_engine.common import recipe from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override @@ -53,6 +55,38 @@ def nvfp4_vanilla(): return nvfp4_recipe +@pytest.mark.parametrize( + "warning_message, expected", + ( + (None, True), + ( + ( + "CUDAGraph.register_generator_state() is deprecated ... " + "It is now a no-op and can be safely removed from your code." + ), + False, + ), + ), +) +def test_graph_rng_registration_capability(monkeypatch, warning_message, expected): + """Distinguish required registration from PyTorch's deprecated no-op.""" + + class FakeGraph: + def register_generator_state(self, generator): + if warning_message is not None: + warnings.warn(warning_message, DeprecationWarning) + + monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) + monkeypatch.setattr(torch.cuda, "default_generators", [object()]) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(graph_module, "graph_safe_rng_available", lambda: True) + graph_module._graph_rng_registration_required.cache_clear() + try: + assert graph_module._graph_rng_registration_required() is expected + finally: + graph_module._graph_rng_registration_required.cache_clear() + + def nvfp4_rht_and_2d_quantization(): nvfp4_recipe = recipe.NVFP4BlockScaling() nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams( diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index be88f730a18..2d2a7fea0ed 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -5,6 +5,7 @@ """Functions for CUDA Graphs support in FP8""" from collections.abc import Iterable import contextlib +import functools import gc import warnings from math import ceil @@ -44,6 +45,19 @@ ) +@functools.lru_cache(maxsize=1) +def _graph_rng_registration_required() -> bool: + """Return whether CUDA graph RNG state must be registered explicitly.""" + if not graph_safe_rng_available(): + return False + graph = torch.cuda.CUDAGraph() + generator = torch.cuda.default_generators[torch.cuda.current_device()] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + graph.register_generator_state(generator) + return not any("is now a no-op" in str(warning.message).lower() for warning in caught) + + class _GraphedCallableHelpers(NamedTuple): """Lifecycle helpers owned by one graphed callable invocation.""" @@ -482,7 +496,7 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params): ) # For cases with multiple active RNG states, e.g. TP. - if graph_safe_rng_available(): + if _graph_rng_registration_required(): for _, state in get_all_rng_states().items(): for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs): fwd_graph.register_generator_state(state) From 3e899ab22942c5984da847b13bb600b51af2babb Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Sun, 8 Nov 2026 10:30:14 -0600 Subject: [PATCH 2/3] Detect CUDA graph RNG behavior directly Signed-off-by: RunguoLi --- tests/pytorch/test_cuda_graphs.py | 49 +++++++++++++++++++---------- transformer_engine/pytorch/graph.py | 21 ++++++++++--- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index b263d1f0b18..f61e716b6f1 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -5,8 +5,8 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import copy +import contextlib import gc -import warnings import weakref import torch @@ -56,29 +56,46 @@ def nvfp4_vanilla(): @pytest.mark.parametrize( - "warning_message, expected", + "advances, expected", ( - (None, True), - ( - ( - "CUDAGraph.register_generator_state() is deprecated ... " - "It is now a no-op and can be safely removed from your code." - ), - False, - ), + (False, True), + (True, False), ), ) -def test_graph_rng_registration_capability(monkeypatch, warning_message, expected): - """Distinguish required registration from PyTorch's deprecated no-op.""" +def test_graph_rng_registration_capability(monkeypatch, advances, expected): + """Detect whether graph replay advances RNG without explicit registration.""" class FakeGraph: - def register_generator_state(self, generator): - if warning_message is not None: - warnings.warn(warning_message, DeprecationWarning) + def __init__(self): + self.replays = 0 + + def replay(self): + self.replays += 1 + + def reset(self): + pass + + class FakeOutput: + def random_(self): + pass + + def clone(self): + return 0 + + class FakeGenerator: + def get_state(self): + return None + + def set_state(self, state): + pass monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) - monkeypatch.setattr(torch.cuda, "default_generators", [object()]) + monkeypatch.setattr(torch.cuda, "default_generators", [FakeGenerator()]) monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "graph", lambda graph: contextlib.nullcontext()) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch, "empty", lambda *args, **kwargs: FakeOutput()) + monkeypatch.setattr(torch, "equal", lambda first, second: not advances) monkeypatch.setattr(graph_module, "graph_safe_rng_available", lambda: True) graph_module._graph_rng_registration_required.cache_clear() try: diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 2d2a7fea0ed..92f1aa9ea0d 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -50,12 +50,23 @@ def _graph_rng_registration_required() -> bool: """Return whether CUDA graph RNG state must be registered explicitly.""" if not graph_safe_rng_available(): return False - graph = torch.cuda.CUDAGraph() + generator = torch.cuda.default_generators[torch.cuda.current_device()] - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - graph.register_generator_state(generator) - return not any("is now a no-op" in str(warning.message).lower() for warning in caught) + original_state = generator.get_state() + graph = torch.cuda.CUDAGraph() + output = torch.empty(8, dtype=torch.int64, device="cuda") + try: + with torch.cuda.graph(graph): + output.random_() + graph.replay() + torch.cuda.synchronize() + first = output.clone() + graph.replay() + torch.cuda.synchronize() + return torch.equal(first, output) + finally: + graph.reset() + generator.set_state(original_state) class _GraphedCallableHelpers(NamedTuple): From a7b0dc3ba7e0fd13a59b3fd937c085ed0e052b1b Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Sun, 8 Nov 2026 10:32:12 -0600 Subject: [PATCH 3/3] Restore graph RNG state on probe failure Signed-off-by: RunguoLi --- transformer_engine/pytorch/graph.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 92f1aa9ea0d..389f064ad51 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -65,8 +65,10 @@ def _graph_rng_registration_required() -> bool: torch.cuda.synchronize() return torch.equal(first, output) finally: - graph.reset() - generator.set_state(original_state) + try: + graph.reset() + finally: + generator.set_state(original_state) class _GraphedCallableHelpers(NamedTuple):