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

[Bug] CUDA graph backward replay illegal memory access after validation with mixed NVFP4/MXFP8 Mamba

Abierto
  • ¿Algo no está bien?
#3,316 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
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
35/100
Tipo de issue
Error
Claridad
Bastante claro
Estado de actividad
Tranquilo
Stack tecnológico
python, pytorch
Área
backend, machine-learning, performance

Línea de trabajo

Comienza con transformer_engine/pytorch/graph.py: backward y reproduce la transición de validación a entrenamiento de iteration-100 usando la configuración de GB200 indicada. Compara el Mamba CUDA graph que falla y la configuración mixta de NVFP4/MXFP8/BF16 con los controles que eliminan Mamba o usan E4M3 MXFP8; se considera terminado cuando el entrenamiento se reanuda después de la validación sin un illegal-memory-access ni un segmentation fault.

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

Descripción

Describe the bug

Nemotron 3 Super training on GB200 consistently hits an illegal memory access around the iteration-100 validation boundary when using Transformer Engine CUDA graphs with a mixed per-module NVFP4, MXFP8, and BF16 recipe.

  • MXFP8 w/ cuda graphs for attn, mamba, moe_router, moe_preprocess: works
  • NVFP4 w/ cuda graphs for attn, mamba, moe_router, moe_preprocess: hits illegal memory access between iterations 90 and 100
  • NVFP4 w/ cuda graphs for attn, moe_router, moe_preprocess: works

The symptom initially appeared between the iteration 90 and 100 log points, but detailed timestamps show that iteration 100 and its 14 validation iterations both complete successfully. The failure occurs during the first training backward after validation. One rank reports a CUDA illegal-memory-access error from the Megatron Core backward call, while another rank can simultaneously segfault inside Transformer Engine CUDA graph backward replay. The failing ranks vary across executions.

Representative stacks:

train_step
  -> forward_backward_no_pipelining
  -> backward_step
  -> custom_backward
  -> Variable._execution_engine.run_backward
torch.AcceleratorError: CUDA error: an illegal memory access was encountered

transformer_engine/pytorch/graph.py: backward
  -> torch/cuda/graphs.py: replay
  -> at::cuda::CUDAGraph::replay
  -> cudaGraphLaunch
  -> cuGraphLaunch
Fatal Python error: Segmentation fault

Steps/Code to reproduce bug

Run Nemotron 3 Super training with this configuration:

Hardware: 16 GB200 nodes, 4 GPUs per node, 64 ranks
Parallelism: TP=2, PP=1, EP=64, ETP=1
Sequence length: 8192
Micro/global batch size: 1/512
Transformer implementation: transformer_engine
Attention backend: fused
CUDA graph implementation: transformer_engine
CUDA graph modules: [attn, mamba, moe_router, moe_preprocess]
CUDA graph warmup steps: 3
MoE dispatcher: flex with HybridEP, 32 SMs
Evaluation interval/iterations: 100/14
Manual GC interval: 100

The base quantization recipe is NVFP4 E2M1. The per-module precision configuration keeps attention QKV/projection, latent projections, and MTP in BF16, and uses MXFP8 for the Mamba output projection. The last 14 layers are also BF16.

The exact mixed-precision arguments used by the failing NVFP4 job were:

--bf16
--grad-reduce-in-bf16
--te-precision-config-file /mnt/artifacts/model/nemotron3_super_release_gb200/te_quant.cfg
--first-last-layers-bf16
--num-layers-at-start-in-bf16 0
--num-layers-at-end-in-bf16 14
--fp4-format e2m1
--fp4-recipe nvfp4

FP4 parameter gather was not enabled (fp4_param_gather=False).

The relevant runtime settings include:

CUDA_DEVICE_MAX_CONNECTIONS=32
NCCL_GRAPH_REGISTER=0
NCCL_NVLS_ENABLE=0
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=64
NUM_OF_TOKENS_PER_CHUNK_COMBINE_API=128
NVTE_FWD_LAYERNORM_SM_MARGIN=20
NVTE_BWD_LAYERNORM_SM_MARGIN=20

Train through iteration 100, run validation, and resume training. In failing runs, iteration 100 completes, validation takes about 112 seconds and completes, then the first resumed training backward faults within several seconds.

What we tested and how we narrowed it down

  • The failure reproduced across multiple independent executions and retries on different nodes and ranks.
  • It reproduced with Transformer Engine 2.18 at two revisions, including 2.18.0+e0a871e6 with the merged weight-preswizzling stale-pointer fix and the later single-grouped-weight fixes present. Those fixes did not resolve this end-to-end failure.
  • It reproduced with both NVSHMEM 3.6.5 and 3.4.5, so the NVSHMEM upgrade is not the cause.
  • We aligned CUDA, NCCL, allocator, layernorm-margin, HybridEP-domain, and token-chunk settings with known working GB200 recipes. The failure remained.
  • A Transformer Engine 2.17 plus cuDNN 9.21 diagnostic completed only five iterations. It did not exercise the iteration-100 boundary, so it is not evidence for or against a version regression.
  • Keeping the updated configuration but removing only Mamba from the CUDA graph module list allowed validation at iteration 100 to complete and training to continue cleanly through iteration 120, with zero skipped or NaN iterations.
  • Keeping Mamba in the full CUDA graph module list but replacing the mixed per-module NVFP4/MXFP8/BF16 recipe with an E4M3 MXFP8 recipe plus BF16 boundary layers allowed validation at iteration 100 to complete and training to continue cleanly through iteration 170, again with zero skipped or NaN iterations.

These controls narrow the problem to an interaction that requires Mamba CUDA graph capture and the mixed per-module NVFP4/MXFP8/BF16 recipe, or state introduced by that recipe across the train-to-validation-to-train transition. It does not look like a generic HybridEP, NVSHMEM, or Mamba training failure.

Expected behavior

Validation should not leave stale or incompatible quantization or CUDA graph state. Training should resume after validation without an illegal memory access during backward graph replay.

Environment overview

  • Environment location: Docker on an internal bare-metal GB200 cluster
  • Transformer Engine install: built from source from the Megatron Core dependency lock during the Docker image build
  • Base image: NVIDIA PyTorch 26.06-derived image

Environment details

  • Python: 3.12
  • PyTorch: 2.13.0a0+8145d630e8.nv26.6
  • Megatron Core: 0.19.0+af9e4408d
  • Transformer Engine: 2.18.0+e0a871e6
  • CUDA: 13.3
  • NCCL: 2.30.7+cuda13.3

Device details

  • GPU model: NVIDIA GB200
  • Topology: 16 nodes, 4 GPUs per node

Additional context

The CUDA error is asynchronous, so the Python error site alone does not identify the captured kernel that first corrupts memory. The simultaneous native stack on another rank localizes the visible failure to Transformer Engine backward graph replay, but we have not yet isolated the exact Mamba or quantization kernel. We can run a launch-blocking diagnostic or help reduce this to a smaller reproducer if there is a preferred instrumentation path.

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

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.