NVIDIA-NeMo / NVIDIA-NeMo/RL

High vllm memory after sleep with tp > 1

Open
#223 2 comments 0 reactions 0 assignees View on GitHub
bug vllm
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

From @yfw : I noticed when setting `policy.generation.vllm_cfg.tensor_parallel_size=2` , vllm doesn't seem to completely free memory when it's supposed to be asleep. It's what's taking that 7-8 gb. It also seemed to have gotten slightly worse with this commit: https://github.com/NVIDIA/reinforcer/commit/e3af33744ded357cc1f52f0e1f04e9f21ede086d
I can see this with llama3.2-1b:
uv run examples/run_grpo_math.py --config examples/configs/grpo_math_1B.yaml grpo.val_at_start=True checkpointing.enabled=False logger.wandb_enabled=False cluster.gpus_per_node=2 cluster.num_nodes=1 policy.precision=bfloat16 policy.model_name=meta-llama/Llama-3.2-1B-Instruct policy.max_total_sequence_length=512 policy.generation.vllm_cfg.tensor_parallel_size=2
This is the vllm memory I'm seeing when it's supposed to be asleep:
Commit https://github.com/NVIDIA/reinforcer/commit/e3af33744ded357cc1f52f0e1f04e9f21ede086d: 7392MiB when asleep
Commit before https://github.com/NVIDIA/reinforcer/commit/dd7c2d757c4f697ede3ffd02173bcc461e0a0d98: 2440MiB
When setting policy.generation.vllm_cfg.tensor_parallel_size=1: (any commit): 774MiB

**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.