oom

Signed-off-by: Chen Zhang <[email protected]>

Author: heheda12345

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()