OOM during NVFP4 quantization of NVIDIA-Nemotron-3-Nano-30B-A3B with Pipeline Parallelism
- 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
Assessment
This issue has not been assessed yet.