NVIDIA-NeMo / NVIDIA-NeMo/RL

GRPO OOM testing script

Open
#884 1 comment 0 reactions 0 assignees View on GitHub
research t-mcore t-pytdensor
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

For testing ooms, right now our current process is to just launch a regular GRPO run and wait for it to oom after a few minutes. this process may take 5-10 mins or more after the job is running depending on the config. it may be helpful to make a grpo oom-testing script which skips the initial vllm cuda graph capture, validation, and rollouts and just runs through the policy train with the configured global/micro batch sizes and seq lens.

A more concrete definition: given an input config, do the regular GRPO setup but do not generate rollouts. In place of the rollouts, create dummy tensors using the max seq len and try to take one train step.

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.