[bug] Qwen3-30B-A3B fp8_mx pretrain OOMs with official defaults on 8x B300, while bf16 and nvfp4 run successfully

[wdkkkkkkkkkk]

Problem

I am running the official Megatron-Bridge performance recipe for Qwen3-30B-A3B inside the official container nvcr.io/nvidia/nemo:26.06.00 on a single node with 8× NVIDIA B300 GPUs.

With the official/default configuration, bf16 and nvfp4 pretraining can run successfully. However, the same official recipe with --compute_dtype fp8_mx consistently runs into CUDA OOM.

I checked the runtime arguments and did not intentionally add custom tuning overrides. The failure seems specific to the fp8_mx path for Qwen3-30B-A3B on B300. I would like to confirm whether the official fp8_mx recipe is expected to run under these default settings, or whether there is a missing required configuration/environment variable for this mode.

Actual behavior: the fp8_mx run fails with CUDA OOM, while bf16 and nvfp4 run successfully under the same environment and recipe.
mxfp8_pretrain_fail.log

Minimal repro

Environment:

Container: nvcr.io/nvidia/nemo:26.06.00
Hardware: single node, 8x NVIDIA B300

Minimal reproduction script:

#!/bin/bash
set -euo pipefail

cd /opt/Megatron-Bridge

export PYTHONPATH=/opt/Megatron-Bridge:/opt/Megatron-Bridge/scripts/performance

NUM_GPUS=8
MAX_STEPS=20

torchrun --standalone --nproc_per_node=$NUM_GPUS \
  scripts/performance/run_script.py \
  --model_family_name qwen \
  --model_recipe_name qwen3_30b_a3b \
  --num_gpus $NUM_GPUS \
  --gpu b300 \
  --compute_dtype fp8_mx \
  --config_variant v1 \
  --data mock \
  --max_steps $MAX_STEPS

Expected behavior

Qwen3-30B-A3B pretraining with --compute_dtype fp8_mx should run successfully under the official/default recipe, similar to the bf16 and nvfp4 variants, and should be able to reproduce the official expected performance level on B300, around 900+ TFLOP/s/GPU, assuming the documented default settings are sufficient.

Affected area

area:training

Regression?

Yes

Environment

Branch / commit:

Repository: [/opt/Megatron-Bridge]

Commit: fcbb603

Container or environment:

Container: nvcr.io/nvidia/nemo:26.06.00
Workdir: /opt/Megatron-Bridge
Launcher: torchrun --standalone --nproc_per_node=8
Script: scripts/performance/run_script.py

GPU / machine:

Single node with 8x NVIDIA B300 SXM6 AC GPUs
Driver Version: 590.48.01
CUDA Version reported by nvidia-smi: 13.1
GPU memory: ~275 GB per GPU
NCCL version observed in log: 2.30.4+cuda13.2

Python / CUDA / PyTorch:

Python: 3.12.3
PyTorch: 2.12.0a0+0291f960b6.nv26.04.48445190
torch.version.cuda: 13.2

Logs

Total number of parameters in billions: 30.52
Number of dense and embedding parameters in most loaded shard in billions: 1.5287
Number of routed expert parameters in most loaded shard in billions: 3.6239
Theoretical memory footprints: weight and optimizer=73142.09 MB
 [2026-06-30 03:45:46] iteration        1/      20 | consumed samples:          512 | elapsed time per iteration (ms): 116509.7 | learning rate: 3.000000E-05 | global batch size:   512 | lm loss: 1.234034E+01 | load_balancing_loss: 1.000110E+00 | loss scale: 1.0 | grad norm: 3.153 | number of skipped iterations:   0 | number of nan iterations:   0 |
[Rank 0] (after 1 iterations) memory (GB) | mem-allocated-gigabytes: 86.177 | mem-active-gigabytes: 86.177 | mem-inactive-gigabytes: 60.88 | mem-reserved-gigabytes: 268.91 | mem-max-allocated-gigabytes: 234.12 | mem-max-active-gigabytes: 276.68 | mem-max-inactive-gigabytes: 88.001 | mem-max-reserved-gigabytes: 283.52 | mem-alloc-retires: 27 | mem-allocated-count: 14188
INFO:megatron.core.transformer.moe.paged_stash:Paged stash: max_tokens_dict is None, skipping stash buffer allocation
[rank7]: Traceback (most recent call last):
[rank7]:   File "/opt/Megatron-Bridge/scripts/performance/run_script.py", line 131, in <module>
[rank7]:     main()
[rank7]:   File "/opt/Megatron-Bridge/scripts/performance/run_script.py", line 123, in main
[rank7]:     pretrain(config=recipe, forward_step_func=forward_step_func)
[rank7]:   File "/opt/Megatron-Bridge/src/megatron/bridge/utils/decorators.py", line 39, in wrapper
[rank7]:     return func(*args, **kwargs)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/pretrain.py", line 98, in pretrain
[rank7]:     _pretrain(state=state, forward_step_func=forward_step_func, callback_manager=callback_manager)
[rank7]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/pretrain.py", line 142, in _pretrain
[rank7]:     train(
[rank7]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/train.py", line 450, in train
[rank7]:     ) = wrapped_train_step(
[rank7]:         ^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/train.py", line 899, in train_step
[rank7]:     losses_reduced = forward_backward_func(
[rank7]:                      ^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/transformer/moe/paged_stash.py", line 1441, in __call__
[rank7]:     result = self.forward_backward_func(*args, **kwargs)
[rank7]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/full_cuda_graph.py", line 230, in __call__
[rank7]:     FullCudaGraphWrapper.result[training_str] = self.forward_backward_func(*args, **kwargs)
[rank7]:                                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 672, in forward_backward_no_pipelining
[rank7]:     forward_data_store, total_num_tokens = combined_1f1b_schedule_for_no_pipelining(
[rank7]:                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/combined_1f1b.py", line 89, in combined_1f1b_schedule_for_no_pipelining
[rank7]:     output_tensor, num_tokens, _ = combined_forward_backward_step(
[rank7]:                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/combined_1f1b.py", line 448, in combined_forward_backward_step
[rank7]:     output_tensor = type(f_schedule_plan or b_schedule_plan).run(
[rank7]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/models/common/model_chunk_schedule_plan.py", line 584, in run
[rank7]:     f_input = f_schedule_plan.post_process.forward(f_input)
[rank7]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/utils.py", line 203, in forward
[rank7]:     return self._forward(*inputs)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/utils.py", line 216, in _forward
[rank7]:     data = self.forward_func(*data)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/models/gpt/fine_grained_callables.py", line 41, in wrapped_func
[rank7]:     return method_ref()(*args, **kwarg)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/models/gpt/fine_grained_callables.py", line 215, in forward_impl
[rank7]:     loss = self.gpt_model._postprocess(
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/models/gpt/gpt_model.py", line 733, in _postprocess
[rank7]:     logits, _ = self.output_layer(
[rank7]:                 ^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank7]:     return self._call_impl(*args, **kwargs)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1884, in _call_impl
[rank7]:     return inner()
[rank7]:            ^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1832, in inner
[rank7]:     result = forward_call(*args, **kwargs)
[rank7]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 1054, in forward
[rank7]:     output_parallel = self._forward_impl(
[rank7]:                       ^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 984, in _forward_impl
[rank7]:     return linear_with_grad_accumulation_and_async_allreduce(input, weight, *args, **kwargs)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 764, in linear_with_grad_accumulation_and_async_allreduce
[rank7]:     return LinearWithGradAccumulationAndAsyncCommunication.apply(*args)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 596, in apply
[rank7]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/dist-packages/torch/amp/autocast_mode.py", line 481, in decorate_fwd
[rank7]:     return fwd(*args, **kwargs)  # pyrefly: ignore [not-callable]
[rank7]:            ^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 501, in forward
[rank7]:     output = torch.matmul(total_input, weight.t())
[rank7]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 9.27 GiB. GPU 7 has a total capacity of 267.69 GiB of which 7.04 GiB is free. Including non-PyTorch memory, this process has 260.63 GiB memory in use. Of the allocated memory 236.15 GiB is allocated by PyTorch, and 17.52 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

[yaoyu-33]

Thanks for the detailed report. I compared the failing log with a successful 8x B300 reference run.

The resolved workload configuration is the same in both cases: EP=8, MBS=8, GBS=512, sequence length 4096, HybridEP, and full-iteration CUDA graphs. The important difference is the launch path: invoking scripts/performance/run_script.py directly bypasses the performance executor and PerfEnvPlugin, which normally inject the allocator and CUDA-graph environment.

Please retry with these variables set before torchrun:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,graph_capture_record_stream_reuse:True
export TORCH_NCCL_AVOID_RECORD_STREAMS=0
export NCCL_NVLS_ENABLE=0
export NCCL_GRAPH_REGISTER=0

The memory data strongly matches allocator fragmentation during graph capture: after iteration 1, the failing run reports about 60.9 GB inactive and 268.9 GB reserved, then falls roughly 2.2 GiB short on the next 9.27 GiB allocation. The reference run has 0 inactive memory at the same point and completes all 50 steps.

These settings are already part of the release's performance launcher defaults and full-iteration graph plugin:

• executor defaults
• full-iteration graph environment

This looks like a direct-launch/documentation gap rather than a different model recipe. Please let us know whether the rerun with these exports completes successfully.

[wdkkkkkkkkkk]

Thanks for the clarification and the detailed explanation.

I tried rerunning the fp8_mx recipe with the environment variables suggested above:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,graph_capture_record_stream_reuse:True
export TORCH_NCCL_AVOID_RECORD_STREAMS=0
export NCCL_NVLS_ENABLE=0
export NCCL_GRAPH_REGISTER=0

However, the run still hit CUDA OOM in my environment.

Following the same direction, I further injected additional performance-related environment variables used by the launcher/plugin path, and with those settings I was able to run the official Qwen3-30B-A3B fp8_mx configuration successfully on 8x B300. The observed performance was around 950+ TFLOP/s/GPU, which is consistent with the expected official-level performance.

export CUDA_DEVICE_MAX_CONNECTIONS=32
export NCCL_IGNORE_CPU_AFFINITY=1
export NVLINK_DOMAIN_SIZE=8
export USE_MNNVL=0
export NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8
export NUM_OF_TOKENS_PER_CHUNK_COMBINE_API=128
export NVTE_FWD_LAYERNORM_SM_MARGIN=20
export NVTE_BWD_LAYERNORM_SM_MARGIN=20
export NVTE_CUTEDSL_FUSED_GROUPED_MLP=1
export CUDNNFE_CLUSTER_OVERLAP_MARGIN=8

This seems to confirm that the issue is not the model recipe itself, but that the direct launch path still misses some required runtime/performance environment setup for this configuration.

I also wanted to ask one follow-up question: for the same Qwen3-30B-A3B recipe on 8x B300, do you have official best-practice configurations and reference performance numbers for the bf16 and nvfp4 precision modes as well? It would be very helpful to compare against the recommended settings and expected throughput for these modes.