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

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.