NVIDIA / NVIDIA/Megatron-LM

OOM during NVFP4 quantization of NVIDIA-Nemotron-3-Nano-30B-A3B with Pipeline Parallelism

Open
#3,311 3 comments 0 reactions 1 assignee Claimed by @ChenhanYu View on GitHub
bug community-request waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

CUDA Out of Memory error occurs during NVFP4 quantization calibration phase for the NVIDIA Nemotron-3-Nano-30B-A3B hybrid Mamba-Transformer model when using Pipeline Parallelism (PP=2). The error happens in the NVFP4 fake quantization kernel when trying to allocate a temporary tensor, despite GPU having sufficient total capacity (2x95GB).

[error.log](https://github.com/user-attachments/files/25174816/error.log)

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB.
GPU 1 has a total capacity of 94.97 GiB of which 20.50 MiB is free.
Including non-PyTorch memory, this process has 94.94 GiB memory in use.
Of the allocated memory 93.10 GiB is allocated by PyTorch, and 58.35 MiB is reserved by PyTorch but unallocated.

**Steps/Code to reproduce bug**

TP=1 PP=2 MLM_SKIP_INSTALL=1 \
HF_MODEL_CKPT=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
MLM_MODEL_SAVE=/tmp/Nemotron-3-Nano-30B-A3B_nvfp4_quant \
MLM_EXTRA_ARGS="--no-gradient-accumulation-fusion --attention-backend unfused" \
./examples/post_training/modelopt/quantize.sh nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 NVFP4_DEFAULT_CFG

**Expected behavior**

Quantization should complete successfully with PP=2, utilizing both GPUs efficiently to fit the model and quantization overhead within the 95GB memory capacity per GPU.

**Additional context**

GPU: 2x RTX PRO 6000
Megatron-LM: `main` branch (commit: a3ec4b02e)
ModelOpt version: `0.42.0rc1.dev9+ge53ca61b7`
Python: 3.12

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.