fix(test): give test_policy_overfit the allocator config its launcher sets - #731
Open
shuangwu wants to merge 1 commit into
Open
fix(test): give test_policy_overfit the allocator config its launcher sets#731shuangwu wants to merge 1 commit into
shuangwu wants to merge 1 commit into
Conversation
… 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
test_policy_overfitOOMs deterministically on 8xH100 80GB and passes onCI'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 cachingallocator. With the default allocator this model leaves the card ~99% held
in retained, fragmented segments, so the next collective fails inside NCCL:
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:Truebacks segments with virtual memory and releasesphysical pages, so NCCL's
cudaMallocsucceeds.cosmos_rl/launcher/utility.pyalready sets exactly this for every replicait starts, which is why real training runs never hit this.
test_policy_overfitspawnstorchrundirectly viasubprocess.Popenandinherits only
os.environ, so it is the one path that misses it. Set it inpolicy_envrather than in the CI harness so the test carries its ownrequirement on any runner;
setdefaultleaves an operator override intact.Measured on one exclusive 8xH100 80GB node, single variable per run:
fsdp_reduce_dtypebf16MixedPrecisionPolicyexpandable_segments:TruePeak 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
vLLMRolloutimport fromtests/launch_test_worker.py. That helper serves 13 modes and only therollout 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
vLLMRolloutin onewith 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.