dlinfer backend support ray (#3903)

* dlinfer backend support ray

* remove ppu change

Author: yao-fengchen

lmdeploy/pytorch/backends/dlinfer/camb/op_backend.py

@@ -123,3 +123,8 @@ def build_graph_runner(model: torch.nn.Module, model_config: ModelConfig, cache_
         """Build graph runner."""
         from lmdeploy.pytorch.backends.cuda.graph_runner import CUDAGraphRunner
         return CUDAGraphRunner(model, model_config, cache_config, backend_config, device)
+
+    @staticmethod
+    def support_ray():
+        """Support ray."""
+        return True

lmdeploy/pytorch/backends/dlinfer/maca/op_backend.py

@@ -109,3 +109,8 @@ def build_graph_runner(model: torch.nn.Module, model_config: ModelConfig, cache_
         """Build graph runner."""
         from lmdeploy.pytorch.backends.cuda.graph_runner import CUDAGraphRunner
         return CUDAGraphRunner(model, model_config, cache_config, backend_config, device)
+
+    @staticmethod
+    def support_ray():
+        """Support ray."""
+        return True

lmdeploy/pytorch/engine/executor/ray_executor.py

@@ -549,11 +549,13 @@ def _init_workers_ray(self, placement_group: PlacementGroup, worker_kwargs: dict
     def _init_distributed_environment_by_device(self, device_str: str):
         """Init distributed environment."""
         driver_ip = _get_master_addr()
-        if device_str == 'cuda':
+        if device_str in ['cuda', 'maca']:
             self.workers = self._sort_workers(driver_ip, self.workers)
 
         elif device_str == 'ascend':
             self._init_ascend_distributed_environment(driver_ip)
+        elif device_str == 'camb':
+            self._init_camb_distributed_environment(driver_ip)
         else:
             raise ValueError(f'Unsupported device type: {device_str}')
 
@@ -576,6 +578,10 @@ def _init_ascend_distributed_environment(self, driver_ip):
         else:
             self.workers = self._sort_workers(driver_ip, self.workers)
 
+    def _init_camb_distributed_environment(self, driver_ip):
+        self.workers = self._sort_workers(driver_ip, self.workers)
+        ray.get([worker.set_device.remote(idx) for idx, worker in enumerate(self.workers)])
+
     """ PD Disaggregation API Begin """
 
     def p2p_initialize(self, init_request: DistServeInitRequest):

lmdeploy/pytorch/ray.py

@@ -16,10 +16,12 @@
 def get_device_str(device_type: str = None) -> str:
     """Get device str."""
     device_type = device_type or get_device_manager().current_context().device_type
-    if device_type == 'cuda':
+    if device_type in ['cuda', 'maca']:
         device_type = 'GPU'
     elif device_type == 'ascend':
         device_type = 'NPU'
+    elif device_type == 'camb':
+        device_type = 'MLU'
     else:
         raise ValueError(f'Unsupported device type: {device_type}')
 

requirements/runtime_camb.txt

@@ -16,7 +16,7 @@ safetensors
 sentencepiece
 shortuuid
 tiktoken
-torch==2.4.0
-torchvision<=0.19.0,>=0.15.0
+torch<=2.6.0,>=2.4.0
+torchvision<=0.21.0,>=0.15.0
 transformers
 uvicorn

requirements/runtime_maca.txt

@@ -16,8 +16,8 @@ safetensors
 sentencepiece
 shortuuid
 tiktoken
-torch<=2.4.0,>=2.0.0
-torchvision<=0.19.0,>=0.15.0
+torch<=2.6.0,>=2.0.0
+torchvision<=0.21.0,>=0.15.0
 transformers
 triton>=2.1.0; sys_platform == "linux"
 uvicorn