oom

Signed-off-by: Chen Zhang <[email protected]>
vllm-project · Jan 8, 2025 · 546e75f · 546e75f
1 parent 2cb84f2
commit 546e75f
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py
@@ -40,6 +40,12 @@ def test_engine_core(monkeypatch):
 
     with monkeypatch.context() as m:
         m.setenv("VLLM_USE_V1", "1")
+        import torch
+        total_allocated_bytes = torch.cuda.mem_get_info(
+        )[1] - torch.cuda.mem_get_info()[0]
+        print("allocated", total_allocated_bytes / (2**30), 'GB', "free",
+              torch.cuda.mem_get_info()[0] / (2**30), 'GB', "total",
+              torch.cuda.mem_get_info()[1] / (2**30), 'GB')
         """Setup the EngineCore."""
         engine_args = EngineArgs(model=MODEL_NAME)
         vllm_config = engine_args.create_engine_config()
@@ -146,6 +152,12 @@ def test_engine_core_advanced_sampling(monkeypatch):
     """
     with monkeypatch.context() as m:
         m.setenv("VLLM_USE_V1", "1")
+        import torch
+        total_allocated_bytes = torch.cuda.mem_get_info(
+        )[1] - torch.cuda.mem_get_info()[0]
+        print("allocated", total_allocated_bytes / (2**30), 'GB', "free",
+              torch.cuda.mem_get_info()[0] / (2**30), 'GB', "total",
+              torch.cuda.mem_get_info()[1] / (2**30), 'GB')
         """Setup the EngineCore."""
         engine_args = EngineArgs(model=MODEL_NAME)
         vllm_config = engine_args.create_engine_config()