NVIDIA-NeMo / NVIDIA-NeMo/RL

[FSDP2] set expandable_segments:True to reduce memory fragmentation

Open
#919 4 comments 0 reactions 1 assignee Claimed by @joyang-nv View on GitHub
bug external t-pytdensor x-pyt
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

We recently verified fsdp2 + Qwen3-30B-A3B can run on a single node with 8 GPUs (80G). To avoid oom, we need `expandable_segments:True` to address memory fragmentation

not sure what's https://github.com/NVIDIA-NeMo/RL/issues/916 about. but hoepfully we can achieve the same thing in nemo-RL
* option 1: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True for the entire program
* option 2: if vllm or sglang complains, we can call `torch.cuda.memory._set_allocator_settings(f"expandable_segments:false")` before wake up and set to true before sleep https://github.com/volcengine/verl/pull/3020#issuecomment-3181828585

Image

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Environment overview (please complete the following information)**

- Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
- Method of install: [pip install or from source]. Please specify exact commands you used to install.
- If method of install is [Docker], provide `docker pull` & `docker run` commands used

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

Add any other context about the problem here.
Example: GPU model

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.