diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 46aab13075..f61e716b6f 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -5,6 +5,7 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import copy +import contextlib import gc import weakref @@ -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,55 @@ def nvfp4_vanilla(): return nvfp4_recipe +@pytest.mark.parametrize( + "advances, expected", + ( + (False, True), + (True, False), + ), +) +def test_graph_rng_registration_capability(monkeypatch, advances, expected): + """Detect whether graph replay advances RNG without explicit registration.""" + + class FakeGraph: + 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", [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: + 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 be88f730a1..389f064ad5 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,32 @@ ) +@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 + + generator = torch.cuda.default_generators[torch.cuda.current_device()] + 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: + try: + graph.reset() + finally: + generator.set_state(original_state) + + class _GraphedCallableHelpers(NamedTuple): """Lifecycle helpers owned by one graphed callable invocation.""" @@ -482,7 +509,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)