[diffusion] fix: bind each rank to accelerator before distributed init (#33054)

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -559,6 +559,7 @@ def maybe_init_distributed_environment_and_model_parallel(
         main_process_only=False,
     )
 
+    current_platform.set_device(device)
     init_distributed_environment(
         world_size=world_size,
         rank=rank,
@@ -577,14 +578,6 @@ def maybe_init_distributed_environment_and_model_parallel(
         sequence_parallel_degree=sp_size,
     )
 
-    # Only set CUDA device if we're on a CUDA platform
-    if current_platform.is_cuda_alike():
-        device = torch.device(f"cuda:{local_rank}")
-        torch.cuda.set_device(device)
-    elif current_platform.is_npu():
-        device = torch.device(f"npu:{local_rank}")
-        torch.npu.set_device(device)
-
 
 def model_parallel_is_initialized() -> bool:
     """Check if model parallel groups are initialized."""

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -222,7 +222,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
 
     def init_device_and_model(self) -> None:
         """Initialize the device and load the model."""
-        torch.get_device_module().set_device(self.local_rank)
+        current_platform.set_device(current_platform.get_device(self.local_rank))
         intra_op_threads = _worker_cpu_intra_op_threads(self.server_args.num_gpus)
         if intra_op_threads is not None:
             torch.set_num_threads(intra_op_threads)

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -225,6 +225,10 @@ class Platform:
     def get_local_torch_device(cls) -> torch.device:
         raise NotImplementedError
 
+    @classmethod
+    def set_device(cls, device: torch.device) -> None:
+        torch.get_device_module(device).set_device(device)
+
     @classmethod
     def get_attn_backend_cls_str(
         cls,

python/sglang/multimodal_gen/runtime/platforms/mps.py

@@ -40,6 +40,10 @@ class MpsPlatform(Platform):
     def get_local_torch_device(cls) -> torch.device:
         return torch.device("mps")
 
+    @classmethod
+    def set_device(cls, device: torch.device) -> None:
+        pass
+
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
         raise NotImplementedError