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

tests/launch_test_worker.py

@@ -35,8 +35,8 @@
 from cosmos_rl.rollout.worker.rollout_control import (
     DisaggregatedRolloutControlWorker,
 )
-from cosmos_rl.rollout.vllm_rollout.vllm_rollout import vLLMRollout
 from cosmos_rl.rollout import State
+
 import types
 from cosmos_rl.dispatcher.command import (
     PolicyToRolloutUnicastCommand,
@@ -434,6 +434,12 @@ def __init__(
             DisaggregatedRolloutControlWorker.consume_command, self
         )
 
+        # Imported here, not at module scope: this helper serves 13 modes and
+        # only the rollout ones need vLLM.  A top-level import made every mode
+        # -- including pure policy/SFT ones that never construct a rollout --
+        # hard-require it, so an image built without vLLM could not run them.
+        from cosmos_rl.rollout.vllm_rollout.vllm_rollout import vLLMRollout
+
         self.rollout = vLLMRollout(self.config, None, torch.cuda.current_device())
 
         self.temp_recv_tensor_queue = Queue()
@@ -1135,6 +1141,8 @@ def rollout_generation(
 
         self.rollout_generation = types.MethodType(rollout_generation, self)
 
+    from cosmos_rl.rollout.vllm_rollout.vllm_rollout import vLLMRollout
+
     vLLMRollout.__init__ = dummy_init
     assert args is not None
     run_rollout(args=args)
@@ -1214,6 +1222,8 @@ def run_rollout_parallelism_extract(rank, fsdp, tp, pp):
         config.policy.model_name_or_path,
         trust_remote_code=True,
     )
+    from cosmos_rl.rollout.vllm_rollout.vllm_rollout import vLLMRollout
+
     rollout = vLLMRollout(config, None, torch.cuda.current_device())
 
     rollout.init_engine(seed=config.rollout.seed, load_format="dummy")

tests/test_policy_overfit.py

@@ -77,6 +77,19 @@ def test_policy_overfit(self):
         ]
         policy_env = dict(os.environ)
         policy_env["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+        # NCCL allocates its buffers with raw cudaMalloc, outside torch's
+        # caching allocator.  With the default allocator this model fills the
+        # card with retained, fragmented segments (~99% of an 80GB H100), and
+        # the next NCCL collective fails in include/alloc.h with
+        # "Cuda failure 2 'out of memory'" -- never a torch "Tried to allocate",
+        # because torch is the one holding it.  Expandable segments release
+        # physical pages back, so NCCL can allocate.
+        #
+        # cosmos_rl/launcher/utility.py sets this for every replica it starts,
+        # which is why real runs are unaffected; this test spawns torchrun
+        # directly and so is the one path that misses it.  Set it here rather
+        # than in the CI harness so the test carries its own requirement.
+        policy_env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
         # Start the process
         policy_process = subprocess.Popen(
             policy_cmd,