[PyTorch] CUDA graph RNG registration floods training logs on automatic-registration builds

Start with transformer_engine/pytorch/graph.py, especially the registration loop around the linked line 484, and compare its capability check with the PyTorch implementations cited in the issue. Then trace the graph setup and existing CUDA graph RNG tests; add coverage for automatic versus required registration, multiple RNG states, and advancement across replay. Done means preserving correct capture/replay while avoiding deprecated no-op calls and their warning flood.

Describe the bug

TE's CUDA graph setup calls CUDAGraph.register_generator_state() once per active RNG tracker state for every forward, backward, and wgrad graph. On the affected NVIDIA PyTorch 26.08 build, this API is a deprecated no-op that emits a C++ warning on every call because registration now happens automatically.

This produces a warning flood during distributed training. With pipeline parallelism, Megatron Core supplies capture samples per layer and microbatch, multiplying the calls. The investigation in NVIDIA-NeMo/Megatron-Bridge#6336 reports approximately 8K–26K warnings per rank for 256-GPU DeepSeek V3 and Qwen3 235B bf16 runs. Some ranks stall writing stderr while others finish capture and wait in collectives; the jobs then hit the 600-second NCCL watchdog timeout after their CUDA graph warmup iterations.

Steps/Code to reproduce bug

The affected path is:

Bridge training → MCore TECudaGraphHelper.create_cudagraphs() → TE make_graphed_callables() → _make_graphed_callables() → register_generator_state().

The registration loop in the Bridge-pinned TE revision registers all three graph types whenever graph_safe_rng_available() is true. That capability check only checks method availability; it does not distinguish required registration from a deprecated no-op.

On the affected PyTorch build, the underlying warning can be demonstrated independently of a full training run:

import torch

generator = torch.Generator(device="cuda")
graph = torch.cuda.CUDAGraph()
for _ in range(3):
    graph.register_generator_state(generator)

This is a minimal illustration of the repeated-warning trigger, not a standalone reproduction of the distributed timeout. The full affected workload uses TE graphs, active RNG tracker states, and pipeline-parallel layer/microbatch capture; the four recipe names and training context are listed in NVIDIA-NeMo/Megatron-Bridge#6336.

The PyTorch implementation at 4fdf77b940 confirms the per-call TORCH_WARN_DEPRECATION and no-op behavior.

Expected behavior / requested TE fix

Please make TE's internal RNG registration conditional on the installed PyTorch implementation's requirements:

• Skip explicit registration when PyTorch handles it automatically.
• Preserve explicit registration on builds where it is required for correct capture/replay.
• Cover both behaviors with tests, including multiple RNG states and RNG advancement across replay, without requiring callers to suppress unrelated C++ warnings.

This should be handled centrally in TE's graph setup so all callers benefit. The inspected make_graphed_callables() API has no registration-policy argument through which MCore could control the inner loop.

Environment overview / device details

• Affected environment: NVIDIA PyTorch 26.08-based training container; reported torch build 2.14.0a0+4fdf77b940.nv26.08.
• Workloads: 256-GPU DeepSeek V3 and Qwen3 235B-A22B bf16 training on GB200/GB300 with TE CUDA graphs.
• Bridge revision for the workaround: 93915172fa1711cca9d3a7fa0721fb15114cd8e8.
• Bridge's TE source pin: 131157b5750b8578cb8efb78554f6fe6a671cc53 (2.20.0 source version).
• The same registration loop is also present in inspected TE v2.18 (27486e03), release_v2.20 (6ea2a74a), and main (5e994640, 2.21.0.dev0). The latter revisions were source-inspected, not independently reproduced on GPUs for this report.

Additional context: workaround and compatibility

NVIDIA-NeMo/Megatron-Bridge#6336 sets TORCH_CPP_LOG_LEVEL=ERROR in four affected recipes before the training interpreter starts. This is a workaround: it suppresses the flood but leaves the unnecessary calls in place, hides unrelated C++ warnings, and does not cover other TE graph callers. We need a TE fix so the global warning suppression can be removed.

Please do not use a blanket torch >= 2.14 skip. NVIDIA/Megatron-LM#7725 reports that the NVIDIA-packaged NGC 26.09 build still requires explicit registration and fails capture if it is skipped, even though the upstream base source also contains the deprecated no-op. Both builds advertise PyTorch 2.14 prereleases. That PR proposes a conservative exact-build exception for the verified 26.08 build; the vendor patch responsible for the packaged-source discrepancy remains unconfirmed. A reliable capability contract would be preferable to a version-only decision.