[Platform] Replace torch.cuda.mem_get_info with torch.accelerator.get_memory_info (#44825)

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>

tests/basic_correctness/test_mem.py

@@ -23,7 +23,7 @@ def test_python_error():
     error happening from the C++ side.
     """
     allocator = get_mem_allocator_instance()
-    total_bytes = current_platform.mem_get_info()[1]
+    total_bytes = torch.accelerator.get_memory_info()[1]
     alloc_bytes = int(total_bytes * 0.7)
     tensors = []
     with allocator.use_memory_pool():
@@ -64,9 +64,9 @@ def test_basic_cumem():
     output = x + y + z
     assert torch.allclose(output, torch.ones_like(output) * 3)
 
-    free_bytes = current_platform.mem_get_info()[0]
+    free_bytes = torch.accelerator.get_memory_info()[0]
     allocator.sleep()
-    free_bytes_after_sleep = current_platform.mem_get_info()[0]
+    free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
     assert free_bytes_after_sleep > free_bytes
     allocator.wake_up()
 
@@ -99,9 +99,9 @@ def test_cumem_with_cudagraph():
     with torch.cuda.graph(model_graph):
         y = model(x)
 
-    free_bytes = current_platform.mem_get_info()[0]
+    free_bytes = torch.accelerator.get_memory_info()[0]
     allocator.sleep()
-    free_bytes_after_sleep = current_platform.mem_get_info()[0]
+    free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
     assert free_bytes_after_sleep > free_bytes
     allocator.wake_up()
 
@@ -132,7 +132,7 @@ def test_cumem_with_cudagraph():
     ],
 )
 def test_end_to_end(model: str):
-    free, total = current_platform.mem_get_info()
+    free, total = torch.accelerator.get_memory_info()
     used_bytes_baseline = total - free  # in case other process is running
     llm = LLM(model, enable_sleep_mode=True)
     prompt = "How are you?"
@@ -144,7 +144,7 @@ def test_end_to_end(model: str):
     # test sleep level 1 here.
     llm.sleep(level=1)
 
-    free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
+    free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
     # now the memory usage is mostly cudagraph memory pool,
     # and it should be less than the model weights (1B model, 2GiB weights)
@@ -164,7 +164,7 @@ def test_end_to_end(model: str):
     llm.sleep(level=1)
     llm.wake_up(tags=["weights"])
 
-    free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+    free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
 
     # should just reallocate memory for weights (1B model, ~2GiB weights)
@@ -181,7 +181,7 @@ def test_end_to_end(model: str):
 @create_new_process_for_each_test()
 def test_deep_sleep():
     model = "hmellor/tiny-random-LlamaForCausalLM"
-    free, total = current_platform.mem_get_info()
+    free, total = torch.accelerator.get_memory_info()
     used_bytes_baseline = total - free  # in case other process is running
     llm = LLM(model, enable_sleep_mode=True)
     prompt = "How are you?"
@@ -191,13 +191,13 @@ def test_deep_sleep():
     # Put the engine to deep sleep
     llm.sleep(level=2)
 
-    free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
+    free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
     assert used_bytes < 3 * GiB_bytes
 
     llm.wake_up(tags=["weights"])
     llm.collective_rpc("reload_weights")
-    free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+    free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
     assert used_bytes < 4 * GiB_bytes
 
@@ -213,7 +213,7 @@ def test_deep_sleep():
 def test_deep_sleep_async():
     async def test():
         model = "hmellor/tiny-random-LlamaForCausalLM"
-        free, total = current_platform.mem_get_info()
+        free, total = torch.accelerator.get_memory_info()
         used_bytes_baseline = total - free  # in case other process is running
         engine_args = AsyncEngineArgs(
             model=model,
@@ -232,7 +232,7 @@ def test_deep_sleep_async():
 
         await llm.wake_up(tags=["weights"])
         await llm.collective_rpc("reload_weights")
-        free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+        free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
         used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
         assert used_bytes < 4 * GiB_bytes