[bug] Qwen3-30B-A3B GB200 MXFP8 perf preset (full_iteration CUDA graphs + HybridEP) OOMs in first optimizer step on PyTorch >=2.14 / TE >=2.18 stacks #6307

Problem

The GB200 MXFP8 performance preset for Qwen3-30B-A3B (qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config, --config_variant v2) runs out of CUDA memory at iteration 3, in the first optimizer step after full-iteration CUDA-graph capture. The job then hangs until the scheduler timeout instead of exiting.

Preset configuration as resolved at runtime:

• cuda_graph_impl=full_iteration
• moe_flex_dispatcher_backend=hybridep
• EP=8, TP=PP=CP=1
• MBS=4, GBS=512, seq_length=4096
• compute_dtype=fp8_mx
• 8 GPUs (2 nodes x 4 GB200)

This is a regression across the framework stack. Both bracket runs used the same GB200 cluster, driver 610.31, and the same Megatron Bridge / Megatron Core commits:

GPU capacity is 184.35 GiB.

Additional observations:

• Fragmentation is not the cause. Reserved-but-unallocated memory is under 400 MiB in every failing run. The peak requirement itself grew (about +9 GB iter-2 max allocated between the passing and first failing stack).
• Non-PyTorch memory grew on the current stack. It is ~10.1 GiB per GPU on the current stack vs ~6.3 GiB on the first failing stack (CUDA context, NCCL/HybridEP buffers). On GPU 0, three other local ranks each also hold an ~898 MiB context. Rank 2 (no peer contexts) OOMs as well, so those contexts are not the only cause.
• The MLM #7309 backport helps but is not enough. Backporting [TE] Preemptively support TE fused CE overwrite_input=True in MCore. NVIDIA/Megatron-LM#7309 (TE parallel_cross_entropy(..., overwrite_input=True), supported since Minimize the memory usage of the fused cross entropy kernel NVIDIA/TransformerEngine#3273) reduces private pools by ~4.6 GiB, but the OOM still occurs at the same site.
• Two GB200 clusters reproduce it. The current stack fails identically on two different GB200 clusters (drivers 580.178.04 and 615.74).
• It still reproduces on the 2026-09-28, 2026-09-30 and 2026-10-02 nightlies of the current stack: 101.02 GiB private pools, same OOM site.
• The full preset env was applied by bootstrap.py (listed below). This is not the missing-env issue from [bug] Qwen3-30B-A3B fp8_mx pretrain OOMs with official defaults on 8x B300, while bf16 and nvfp4 run successfully #4587.

Minimal repro

# 8 GB200 GPUs (2 nodes x 4), one process per GPU; 8 GPUs is the preset minimum (EP=8)
python scripts/performance/bootstrap.py \
  --model_family_name qwen --model_recipe_name qwen3_30b_a3b \
  --gpu gb200 --compute_dtype fp8_mx --config_variant v2 \
  --domain llm --task pretrain --num_gpus 8 --max_steps 30 \
  --seq_length 4096 --data mock \
  checkpoint.save=null checkpoint.load=null \
  train.eval_interval=30 train.eval_iters=0 logger.log_interval=1

# Preset env applied by bootstrap.py:
# CUDA_DEVICE_MAX_CONNECTIONS=32 CUDNNFE_CLUSTER_OVERLAP_MARGIN=8
# NCCL_GRAPH_REGISTER=0 NCCL_NVLS_ENABLE=0
# NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8 NUM_OF_TOKENS_PER_CHUNK_COMBINE_API=128 NVLINK_DOMAIN_SIZE=72
# NVTE_BWD_LAYERNORM_SM_MARGIN=20 NVTE_FWD_LAYERNORM_SM_MARGIN=20 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1
# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,graph_capture_record_stream_reuse:True
# TORCH_NCCL_AVOID_RECORD_STREAMS=0 TORCH_NCCL_HIGH_PRIORITY=1 USE_MNNVL=1

1. Run the command on a PyTorch 2.13 / TE 2.16.1 / CUDA 13.3 / NCCL 2.30.4 stack: 30 iterations complete.
2. Run it on a PyTorch 2.14 / TE 2.18 / CUDA 13.4.1 / NCCL 2.30.7 stack, or newer: iteration 3 fails with torch.OutOfMemoryError in DistributedOptimizer._copy_model_grads_to_main_grads, and the job hangs afterwards.

Expected behavior

The preset should fit in GB200 memory on current framework stacks and complete all 30 iterations, as it does on the PyTorch 2.13 / TE 2.16.1 stack. The Qwen3-30B-A3B model verification card (examples/model_verification_cards/qwen3-30b-a3b/card.yaml) reports 927.68 TFLOP/s/GPU for this preset on 8x GB200 with the 26.06 NeMo container. If an OOM does occur, the job should exit instead of hanging until the scheduler timeout.

Actual behavior

1. Iterations 1 and 2 complete, including warmup and full-iteration CUDA-graph capture. Iteration 2 reports mem-max-allocated 167.05 GB on the current stack with the #7309 backport (171.19 GB without it), and 175.93 GB on the PyTorch 2.14 stack.
2. In iteration 3, ranks 0, 2, 4 and others raise torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 12.00 MiB in DistributedOptimizer.copy_group_grads (shard_main_param.grad = shard_model_grad.float()). Only a few MiB of the 184.35 GiB are free, and 96.39–105.72 GiB is held in CUDA-graph private pools.
3. The processes then stop making progress but do not exit. The job runs until the scheduler time limit kills it, and no throughput metrics are produced.
4. This happens on every run of the failing stacks: two GB200 clusters and several nightly builds.

Affected area

area:perf

Regression?

Yes

Environment

Branch / commit:

• Passing and first-failing runs: Megatron Bridge 0.6.1+51885cf, Megatron Core 0.19.0+cd4afff.
• Current failing run: Megatron Bridge 0.6.2+c0e164e, Megatron Core 0.19.1+a0f793dfa4e776d99a8aa63bed74e0fafc91b0db.

Container or environment: PyTorch-based framework containers. Exact versions per run are in the table above. deep_ep 58397ea on the current stack. python -m torch.utils.collect_env output was not captured in CI; versions are taken from the job logs.

GPU / machine: GB200 NVL, 4 GPUs/node x 2 nodes, 184.35 GiB per GPU. The bracket used driver 610.31; the current stack was run with drivers 580.178.04 and 615.74.

Python / CUDA / PyTorch: Python 3.12. See the table for CUDA and PyTorch.

Logs

[Rank 0] (after 2 iterations) memory (GB) | mem-allocated-gigabytes: 85.178 | mem-reserved-gigabytes: 165.39 | mem-max-allocated-gigabytes: 167.05 | mem-max-active-gigabytes: 181.14
...
  File "megatron/bridge/training/train.py", line 939, in train_step
  File "megatron/core/optimizer/optimizer.py", line 1862, in step
  File "megatron/core/optimizer/optimizer.py", line 1631, in prepare_grads
  File "megatron/core/optimizer/optimizer.py", line 751, in prepare_grads
    self._copy_model_grads_to_main_grads()
  File "megatron/core/optimizer/distrib_optimizer.py", line 2789, in _copy_model_grads_to_main_grads
    copy_group_grads(self.model_float16_groups, self.shard_fp32_from_float16_groups)
  File "megatron/core/optimizer/distrib_optimizer.py", line 2782, in copy_group_grads
    shard_main_param.grad = shard_model_grad.float()
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 12.00 MiB. GPU 0 has a total capacity of 184.35 GiB of which 13.56 MiB is free. Process ... has 898.00 MiB memory in use. Including non-PyTorch memory, this process has 181.63 GiB memory in use. Process ... has 898.00 MiB memory in use. Process ... has 898.00 MiB memory in use. Of the allocated memory 171.13 GiB is allocated by PyTorch, with 96.39 GiB allocated in private pools (e.g., CUDA Graphs), and 380.24 MiB is reserved by PyTorch but unallocated. expandable_segments is enabled.

Suggested investigation

• Run a one-axis A/B between the passing and failing stacks (swap PyTorch, TE, and CUDA/NCCL one at a time) to find which component grew the peak and the full-iteration graph private pool.
• Check whether the preset's memory budget (MBS=4 with full_iteration graphs) still has headroom for the eager fp32 main-grad materialization in copy_group_grads after capture.
• Check why other local ranks create ~898 MiB contexts on GPU 0, and why non-PyTorch memory grew on the current stack.
• On the current stack, iteration-2 mem-max-active exceeds mem-max-allocated by ~11–14 GB (182.39 vs 171.19 GB without the #7309 backport; 181.14 vs 167.05 GB with it). Check whether frees deferred by record_stream add to the peak. The preset sets TORCH_NCCL_AVOID_RECORD_STREAMS=0 and graph_capture_record_stream_reuse:True.
• Once [Dev] Add combined 1F1B schedule-aware tensor release NVIDIA/Megatron-LM#7062 reaches main and the preset can set it, A/B ep_overlap_use_scheduled_tensor_release=True against the default on this preset. The preset uses overlap_moe_expert_parallel_comm=True with delay_wgrad_compute=True, i.e. the combined 1F1B EP-overlap path.

Related but distinct:

• [bug] Qwen3-30B-A3B fp8_mx pretrain OOMs with official defaults on 8x B300, while bf16 and nvfp4 run successfully #4587: B300 OOM caused by missing preset env on direct launch.
• [BUG] MoE buffers dominate memory fragmentation with MFSDP v2 and CUDA graphs NVIDIA/Megatron-LM#7622: MoE buffer fragmentation with CUDA graphs on GB300 MFSDP v2, CP16, 128K. Its draft fix Replace ShortcutMoE record_stream with explicit synchronization NVIDIA/Megatron-LM#7654 replaces record_stream in the hybrid-model ShortcutMoEBlock. That module is absent at the tested Megatron Core commit and is not used by this GPT MoE + HybridEP path.
• [TE] Preemptively support TE fused CE overwrite_input=True in MCore. NVIDIA/Megatron-LM#7309: tested; partial saving only.
• [Dev] Add combined 1F1B schedule-aware tensor release NVIDIA/Megatron-LM#7062 (merged to dev only): adds the default-off option ep_overlap_use_scheduled_tensor_release for combined 1F1B EP overlap. Its predecessor [Dev] Reduce HybridEP dispatch pending frees with stream hand-back NVIDIA/Megatron-LM#6883 reported about 720 MiB savings, well below the ~9 GB peak growth. Not tested here.

Patch verification

Candidate 1: Does not fix (code patch)

Base revision: a0f793dfa4e776d99a8aa63bed74e0fafc91b0db

Scope: Backport of NVIDIA/Megatron-LM#7309 (TE parallel cross-entropy overwrite_input=True) onto MCore a0f793d, to test whether the extra BF16 logits copy explains the full-iteration CUDA-graph OOM. Pass criterion: 30 iterations complete without OOM.

• CUDA-graph private pools drop from 101.02 GiB to 96.39 GiB (~4.6 GiB saved) with the backport.
• OOM still occurs at iteration 3 in DistributedOptimizer copy_group_grads (shard_model_grad.float()); iter-2 max allocated 167.05 GB.
• The same backported build also OOMs identically on a second GB200 cluster (96.39 GiB pools, 171.13 GiB allocated by PyTorch).

Limitations:

• Tested on the current PyTorch 2.15.0a0 / TE 2.20.0 stack only; not a regression-bracket fix.

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 8f0117c68..fbec345a8 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -3485,6 +3485,9 @@ try:
     from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy
 
     _TE_SUPPORTS_CG_CAPTURABLE = is_te_min_version("2.7.0")
+    _TE_FUSED_PARALLEL_CE_OVERWRITE_INPUT = (
+        "overwrite_input" in inspect.signature(parallel_cross_entropy).parameters
+    )
     current_te_version = get_te_version()
 
     def te_parallel_cross_entropy(
@@ -3492,15 +3495,24 @@ try:
         labels: torch.Tensor,
         tp_group: torch.distributed.ProcessGroup,
         is_cg_capturable: bool = False,
+        overwrite_input: bool = True,
     ):
         """Wrapper function for TE's Cross Entropy Loss kernel"""
+        parallel_cross_entropy_kwargs = {
+            "label_smoothing": 0.0,
+            "reduce_loss": False,
+            "dist_process_group": tp_group,
+        }
+        if _TE_FUSED_PARALLEL_CE_OVERWRITE_INPUT:
+            # TransformerEngine will reuse the input buffer for dgrad if overwrite_input=True.
+            # Supported after https://github.com/NVIDIA/TransformerEngine/pull/3273.
+            parallel_cross_entropy_kwargs["overwrite_input"] = overwrite_input
         if _TE_SUPPORTS_CG_CAPTURABLE:
+            # Use the CUDA graph-capturable version of the loss function.
+            parallel_cross_entropy_kwargs["is_cg_capturable"] = is_cg_capturable
             # According to TE CrossEntropyFunction, ignore_idx defaults to -100
-            return parallel_cross_entropy(
-                logits, labels, 0.0, False, tp_group, -100, is_cg_capturable
-            )
-        else:
-            return parallel_cross_entropy(logits, labels, 0.0, False, tp_group)
+            parallel_cross_entropy_kwargs["ignore_idx"] = -100
+        return parallel_cross_entropy(logits, labels, **parallel_cross_entropy_kwargs)
 
 except ImportError:
     te_parallel_cross_entropy = None  # type: ignore[assignment, misc]

This issue was drafted with assistance from the opus AI model.