NVIDIA-NeMo / NVIDIA-NeMo/RL

Out of Memory Error When Validation and Checkpointing Periods Are Misaligned

Open
#1,137 4 comments 0 reactions 0 assignees View on GitHub
bug external t-checkpoint x-domyn
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

NeMo RL encounters a CUDA out of memory error when val_period and checkpointing.save_period are configured with different values, specifically when checkpointing occurs more frequently than validation (I didn't try different configurations)

**Steps/Code to reproduce bug**

Working Configuration

```yaml
grpo:
val_period: 20
...
checkpointing:
save_period: 20
```
This configuration runs without memory issues.

```yaml
grpo:
val_period: 20
...
checkpointing:
save_period: 10
```
This configuration causes CUDA out of memory errors at step 21 (immediately after the second validation at step 20).

**Timeline Analysis**
From the logs, the sequence of events is:

Step 10: Checkpointing occurs, memory is managed properly
Steps 11-19: Training continues normally
Step 20: Validation runs
Step 20: Another checkpoint is saved after validation
Step 21: Attempt to start next batch generation fails with OOM

Here are highlights of the full log:

```bash
...

========================= Step 10/10078 =========================
▶ Preparing batch...
▶ Generating responses for batch of size 32...
(VllmGenerationWorker pid=1101, ip=100.64.0.14) INFO 09-16 09:53:11 [block_pool.py:321] Successfully reset prefix cache [repeated 7x across cluster]
(MegatronPolicyWorker[rank=8] pid=4950, ip=100.64.0.14) GPU Memory before optimizer offload: 28.74GB allocated, 28.76GB reserved [repeated 33x across cluster]
(RayWorkerWrapper pid=3032, ip=100.64.0.9) INFO 09-16 09:53:13 [gpu_worker.py:104] Sleep mode freed 124.54 GiB memory, 7.95 GiB memory is still in use. [repeated 31x across cluster]
(VllmGenerationWorker pid=2243, ip=100.64.0.9) INFO 09-16 09:53:13 [executor_base.py:187] It took 1.904843 seconds to fall asleep. [repeated 3x across cluster]
(MegatronPolicyWorker[rank=4] pid=18088) GPU Memory after optimizer offload: 4.43GB allocated, 4.53GB reserved [repeated 31x across cluster]
(MegatronPolicyWorker[rank=4] pid=18088) GPU Memory before optimizer offload: 30.77GB allocated, 30.79GB reserved [repeated 31x across cluster]
(MegatronPolicyWorker[rank=16] pid=5243, ip=100.64.0.9) GPU Memory after optimizer offload: 4.14GB allocated, 4.30GB reserved
(MegatronPolicyWorker[rank=22] pid=5536, ip=100.64.0.9) GPU Memory after optimizer offload: 4.14GB allocated, 4.21GB reserved
(MegatronPolicyWorker[rank=20] pid=5538, ip=100.64.0.9) GPU Memory after optimizer offload: 4.14GB allocated, 4.30GB reserved [repeated 30x across cluster]
(VllmGenerationWorker pid=14172) INFO 09-16 09:54:05 [executor_base.py:203] It took 0.724005 seconds to wake up tags ['weights'].
[Refit] Split 563 keys into 7 groups
Adding requests: 100%|██████████| 8/8 [00:00<00:00, 10436.84it/s]
Processed prompts: 0%| | 0/8 [00:00)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/worker/worker_base.py", line 620, in execute_method
raise e
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/worker/worker_base.py", line 611, in execute_method
return run_method(self, method, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/utils/__init__.py", line 2985, in run_method
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/v1/worker/gpu_worker.py", line 113, in wake_up
allocator.wake_up(tags)
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/device_allocator/cumem.py", line 225, in wake_up
create_and_map(handle)
File "/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/lib/python3.12/site-packages/vllm/device_allocator/cumem.py", line 78, in create_and_map
python_create_and_map(*allocation_handle)
RuntimeError: CUDA Error: out of memory at /workspace/csrc/cumem_allocator.cpp:62
(RayWorkerWrapper pid=1965, ip=100.64.0.14) CUDA Error: out of memory at /workspace/csrc/cumem_allocator.cpp:62
```

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

Environment location: Lepton running a NemoRL image built as [stated here](https://docs.nvidia.com/nemo/rl/latest/docker.html#build-docker-images).

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.