Skip to content

fix(test): give test_policy_overfit the allocator config its launcher sets - #731

Open
shuangwu wants to merge 1 commit into
nvidia-cosmos:mainfrom
shuangwu:fix/test-policy-overfit-allocator-config
Open

fix(test): give test_policy_overfit the allocator config its launcher sets#731
shuangwu wants to merge 1 commit into
nvidia-cosmos:mainfrom
shuangwu:fix/test-policy-overfit-allocator-config

Conversation

@shuangwu

Copy link
Copy Markdown
Collaborator

test_policy_overfit OOMs deterministically on 8xH100 80GB and passes on
CI's larger cards. It is not a capacity problem, and the failure is not
torch's.

NCCL allocates its buffers with a raw cudaMalloc, outside torch's caching
allocator. With the default allocator this model leaves the card ~99% held
in retained, fragmented segments, so the next collective fails inside NCCL:

include/alloc.h:228 NCCL WARN Cuda failure 2 'out of memory'
ncclUnhandledCudaError: Call to CUDA function failed.

The tell is what is absent: no torch "Tried to allocate ... GiB" anywhere.
torch is not failing to allocate -- torch is the one holding the memory, and
NCCL is the caller that finds nothing left.

expandable_segments:True backs segments with virtual memory and releases
physical pages, so NCCL's cudaMalloc succeeds.
cosmos_rl/launcher/utility.py already sets exactly this for every replica
it starts, which is why real training runs never hit this.
test_policy_overfit spawns torchrun directly via subprocess.Popen and
inherits only os.environ, so it is the one path that misses it. Set it in
policy_env rather than in the CI harness so the test carries its own
requirement on any runner; setdefault leaves an operator override intact.

Measured on one exclusive 8xH100 80GB node, single variable per run:

run peak / 81,559 MiB result
baseline 80,847 OOM
fsdp_reduce_dtype bf16 80,723 OOM
bf16 + FSDP MixedPrecisionPolicy 80,724 OOM
expandable_segments:True 81,064 30/30

Peak memory went UP and the test passed. It was never about using less
memory, which is why every dtype lever moved it by ~1 MiB: they shrink live
tensors, while the failure is one allocator starving another. Loss tracks
CI's run (13.47852 -> 13.47676 over 30 steps), so numerics are unchanged.

Also drop an unconditional vLLMRollout import from
tests/launch_test_worker.py. That helper serves 13 modes and only the
rollout ones construct a rollout, but the module-level import made every
mode -- including the policy/SFT ones this test uses -- hard-require vLLM.
Imported at the three call sites instead. Verified both ways: the module now
imports in an image without vLLM, and still resolves vLLMRollout in one
with it.

Ruled out along the way, each by measurement rather than argument: node
contention (--exclusive, nvidia-smi clean); cross-suite accumulation
(reproduces in isolation); GPU count (the test hardcodes world_size = 8);
/dev/shm (1.8T) and locked memory (unlimited); NCCL_NVLS_ENABLE
(226 MiB, both arms OOM); and the FSDP reduction dtype above.

… sets

`test_policy_overfit` OOMs deterministically on 8xH100 80GB and passes on
CI's larger cards. It is not a capacity problem, and the failure is not
torch's.

NCCL allocates its buffers with a raw `cudaMalloc`, outside torch's caching
allocator. With the default allocator this model leaves the card ~99% held
in retained, fragmented segments, so the next collective fails inside NCCL:

```text
include/alloc.h:228 NCCL WARN Cuda failure 2 'out of memory'
ncclUnhandledCudaError: Call to CUDA function failed.
```

The tell is what is *absent*: no torch "Tried to allocate ... GiB" anywhere.
torch is not failing to allocate -- torch is the one holding the memory, and
NCCL is the caller that finds nothing left.

`expandable_segments:True` backs segments with virtual memory and releases
physical pages, so NCCL's `cudaMalloc` succeeds.
`cosmos_rl/launcher/utility.py` already sets exactly this for every replica
it starts, which is why real training runs never hit this.
`test_policy_overfit` spawns `torchrun` directly via `subprocess.Popen` and
inherits only `os.environ`, so it is the one path that misses it. Set it in
`policy_env` rather than in the CI harness so the test carries its own
requirement on any runner; `setdefault` leaves an operator override intact.

Measured on one exclusive 8xH100 80GB node, single variable per run:

| run                                   | peak / 81,559 MiB | result   |
|---------------------------------------|-------------------|----------|
| baseline                              | 80,847            | OOM      |
| `fsdp_reduce_dtype` bf16              | 80,723            | OOM      |
| bf16 + FSDP `MixedPrecisionPolicy`    | 80,724            | OOM      |
| **`expandable_segments:True`**        | **81,064**        | **30/30**|

Peak memory went UP and the test passed. It was never about using less
memory, which is why every dtype lever moved it by ~1 MiB: they shrink live
tensors, while the failure is one allocator starving another. Loss tracks
CI's run (13.47852 -> 13.47676 over 30 steps), so numerics are unchanged.

Also drop an unconditional `vLLMRollout` import from
`tests/launch_test_worker.py`. That helper serves 13 modes and only the
rollout ones construct a rollout, but the module-level import made every
mode -- including the policy/SFT ones this test uses -- hard-require vLLM.
Imported at the three call sites instead. Verified both ways: the module now
imports in an image without vLLM, and still resolves `vLLMRollout` in one
with it.

Ruled out along the way, each by measurement rather than argument: node
contention (`--exclusive`, nvidia-smi clean); cross-suite accumulation
(reproduces in isolation); GPU count (the test hardcodes `world_size = 8`);
`/dev/shm` (1.8T) and locked memory (unlimited); `NCCL_NVLS_ENABLE`
(226 MiB, both arms OOM); and the FSDP reduction dtype above.
@shuangwu
shuangwu requested a review from lfengad August 31, 2026 20:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant