fix rollout test to 2GPU

volcengine · Dec 7, 2024 · 23cffb4 · 23cffb4
1 parent 7425b36
commit 23cffb4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/rollout/test_vllm_hf_loader.py b/tests/rollout/test_vllm_hf_loader.py
@@ -27,7 +27,7 @@
 
 
 def test_vllm_with_hf():
-    assert torch.cuda.device_count() >= 4, 'At least 4 GPUs is required to run tp+dp tests.'
+    assert torch.cuda.device_count() >= 2, 'At least 2 GPUs is required to run tp+dp tests.'
 
     # fill rollout config
     max_prompt_length = 16