fix some bug

Signed-off-by: noemotiovon <757486878@qq.com>

Author: noemotiovon

roll/platforms/cpu.py

@@ -13,7 +13,7 @@ class CpuPlatform(Platform):
     communication_backend: str = "gloo"
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         pass
 
     @classmethod

roll/platforms/cuda.py

@@ -16,7 +16,7 @@ class CudaPlatform(Platform):
     communication_backend: str = "nccl"
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         torch._C._cuda_clearCublasWorkspaces()
 
     @classmethod
@@ -39,7 +39,7 @@ def get_vllm_worker_class(clas):
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
 
     @classmethod
-    def set_allocator_settings(cls):
+    def set_allocator_settings(cls) -> None:
         torch.cuda.memory._set_allocator_settings("expandable_segments:False")
 
     @classmethod

roll/platforms/npu.py

@@ -14,7 +14,7 @@ class NpuPlatform(Platform):
     communication_backend: str = "hccl"
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         pass
 
     @classmethod
@@ -37,7 +37,7 @@ def get_vllm_worker_class(clas):
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
 
     @classmethod
-    def set_allocator_settings(cls):
+    def set_allocator_settings(cls) -> None:
         pass
 
     @classmethod

roll/platforms/platform.py

@@ -89,7 +89,7 @@ def __getattr__(self, key: str):
             return None
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         raise NotImplementedError
 
     @classmethod
@@ -98,16 +98,33 @@ def get_vllm_worker_class(cls):
         raise NotImplementedError
 
     @classmethod
-    def set_allocator_settings(cls):
+    def set_allocator_settings(cls) -> None:
         """Configure memory allocator settings based on the device type."""
         raise NotImplementedError
 
     @classmethod
     def get_custom_env_vars(cls) -> dict:
+        """
+        Return custom environment variables specific to the platform.
+
+        Returns:
+            dict: A dictionary of environment variable key-value pairs.
+        """
         raise NotImplementedError
 
     @classmethod
-    def update_env_vars_for_visible_devices(cls, env_vars: dict, gpu_ranks: list):
+    def update_env_vars_for_visible_devices(cls, env_vars: dict, gpu_ranks: list) -> None:
+        """
+        Update environment variables to control device visibility.
+
+        Args:
+            env_vars (dict): Dictionary of current environment variables to modify.
+            gpu_ranks (list): List of device IDs to expose to the process.
+
+        Behavior:
+            - Sets the platform-specific visibility environment variable.
+            - Sets the corresponding Ray experimental flag if needed.
+        """
         visible_devices_env_vars = {
             cls.device_control_env_var: ",".join(map(str, gpu_ranks)),
             cls.ray_experimental_noset: "1",
@@ -116,6 +133,12 @@ def update_env_vars_for_visible_devices(cls, env_vars: dict, gpu_ranks: list):
 
     @classmethod
     def get_visible_gpus(cls) -> list:
+        """
+        Return the list of currently visible device IDs.
+
+        Returns:
+            list: A list of device ID strings parsed from the visibility environment variable.
+        """
         if cls.device_control_env_var is not None:
             return os.environ.get(cls.device_control_env_var, "").split(",")
         return []

roll/platforms/rocm.py

@@ -16,7 +16,7 @@ class RocmPlatform(Platform):
     communication_backend: str = "nccl"
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         torch._C._cuda_clearCublasWorkspaces()
 
     @classmethod
@@ -39,7 +39,7 @@ def get_vllm_worker_class(clas):
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
 
     @classmethod
-    def set_allocator_settings(cls):
+    def set_allocator_settings(cls) -> None:
         torch.cuda.memory._set_allocator_settings("expandable_segments:False")
 
     @classmethod

roll/platforms/unknown.py

@@ -16,7 +16,7 @@ class UnknownPlatform(Platform):
     communication_backend: str = "nccl"
 
     @classmethod
-    def clear_cublas_workspaces(cls):
+    def clear_cublas_workspaces(cls) -> None:
         torch._C._cuda_clearCublasWorkspaces()
 
     @classmethod
@@ -39,7 +39,7 @@ def get_vllm_worker_class(clas):
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
 
     @classmethod
-    def set_allocator_settings(cls):
+    def set_allocator_settings(cls) -> None:
         torch.cuda.memory._set_allocator_settings("expandable_segments:False")
 
     @classmethod

roll/third_party/vllm/vllm_0_8_4/llm.py

@@ -216,6 +216,8 @@ def update_parameter(self, parameter_name, weight, ranks_in_worker):
                 "dtype": weight.dtype,
                 "weight": weight.cpu().tolist()
             }
+        else:
+            weight_dict = weight
         self.collective_rpc(method="update_parameter", args=(parameter_name, weight_dict, ranks_in_worker))
 
     def update_parameter_in_bucket(self, meta_infos, buffer, ranks_in_worker):

roll/third_party/vllm/vllm_0_8_4/v1/worker.py

@@ -26,7 +26,7 @@ def __init__(self, *args, **kwargs):
 
     def update_parameter(self, parameter_name, weight, ranks_in_worker):
         weight_dict = weight
-        weight = torch.tensor(weight_dict["weight"], dtype=weight_dict["dtype"]).cuda()
+        weight = torch.tensor(weight_dict["weight"], dtype=weight_dict["dtype"]).to(current_platform.device_type)
         super().update_parameter(parameter_name, weight, ranks_in_worker)
 
     def broadcast_bucket(self, src_pp_rank, meta_infos, bucket_size):
@@ -35,7 +35,7 @@ def broadcast_bucket(self, src_pp_rank, meta_infos, bucket_size):
 
     def update_parameter_in_bucket(self, meta_infos, buffer, ranks_in_worker):
         RecvBucketManager.dict_to_meta(meta_infos)
-        buffer = torch.tensor(buffer, dtype=torch.int8, device='cuda')
+        buffer = torch.tensor(buffer, dtype=torch.int8, device=current_platform.device_type)
         super().update_parameter_in_bucket(meta_infos, buffer, ranks_in_worker)
 
     def add_lora(self, peft_config) -> bool:

tests/models/cuda_mem/utils.py

@@ -1,10 +1,12 @@
 import torch
 
+from roll.platforms import current_platform
+
 
 def log_gpu_memory_usage(head: str):
-    memory_allocated = torch.cuda.memory_allocated() / 1024**3
-    memory_reserved = torch.cuda.memory_reserved() / 1024**2
-    memory_reserved_max = torch.cuda.max_memory_reserved() / 1024**3
+    memory_allocated = current_platform.memory_allocated() / 1024**3
+    memory_reserved = current_platform.memory_reserved() / 1024**2
+    memory_reserved_max = current_platform.max_memory_reserved() / 1024**3
     message = (
         f"{head}, memory allocated (GB): {memory_allocated}, memory reserved (MB): {memory_reserved}, "
         f"memory max reserved (GB): {memory_reserved_max}"