NVIDIA-NeMo / NVIDIA-NeMo/RL

PPO default yaml OOM

Open
#3,793 1 comment 0 reactions 2 assignees Claimed by @bg51717 View on GitHub
bug Memory
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

PPO default yaml OOM for both backend:
1. DTensor: `examples/configs/ppo_math_1B.yaml`
2. Megatron: `examples/configs/ppo_math_1B_megatron.yaml`
3. SC + Megatron: `examples/configs/ppo_math_1B_megatron_single_controller.yaml` (will be added after #3773)

run with `uv run python examples/run_ppo.py --config examples/configs/ppo_math_1B.yaml`.

error log:
```
ray.exceptions.RayTaskError(OutOfMemoryError): ray::DTensorPolicyWorkerV2.get_logprobs() (pid=2671339, ip=10.65.16.209, actor_id=d5836b60449efbd220b3c2f406000000, repr=DTensorPolicyWorkerV2[rank=0])
File "/root/.local/share/uv/python/cpython-3.13.14-linux-x86_64-gnu/lib/python3.13/concurrent/futures/_base.py", line 453, in result
return self.__get_result()
~~~~~~~~~~~~~~~~~^^
File "/root/.local/share/uv/python/cpython-3.13.14-linux-x86_64-gnu/lib/python3.13/concurrent/futures/_base.py", line 402, in __get_result
raise self._exception
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/utils/nsys.py", line 119, in wrapper
ret = func(*args, **kwargs)
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py", line 653, in get_logprobs
token_logprobs, _metrics, _ = forward_with_post_processing_fn(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
model=self.model,
^^^^^^^^^^^^^^^^^
...<5 lines>...
sequence_dim=sequence_dim,
^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 401, in forward_with_post_processing_fn
result = post_processing_fn(
logits=logits,
...<4 lines>...
sequence_dim=sequence_dim,
)
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 755, in __call__
token_logprobs = self._compute_local_logprobs(
logits, processed_inputs.input_ids
)
File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 838, in _compute_local_logprobs
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/functional.py", line 2267, in log_softmax
ret = input.log_softmax(dim)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 37.09 GiB. GPU 0 has a total capacity of 79.11 GiB of which 10.58 GiB is free. Process 2669658 has 1.35 GiB memory in use. Including non-PyTorch memory, this process has 66.43 GiB memory in use. Process 2672344 has 720.00 MiB memory in use. Of the allocated memory 64.36 GiB is allocated by PyTorch, and 875.71 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
```

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.