Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

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

Abierto
  • ¿Algo no está bien?
#3,645 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
50/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
performance

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3.6k
Forks
851
Merge medio
4 d 15 h
PR fusionados (30 d)
51

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.