[diffusion] hardware: support diffusion (single GPU, 3/N) (#17105)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
2026-01-16 17:01:09 +08:00
parent d9ed80b9f1
commit a1dd3d48ac
13 changed files with 425 additions and 33 deletions
@@ -242,8 +242,12 @@ def init_distributed_environment(
            "distributed environment"
        )

-        # For MPS, don't pass device_id as it doesn't support device indices
-        extra_args = {} if current_platform.is_mps() else dict(device_id=device_id)
+        # For MPS and MUSA, don't pass device_id as it doesn't support device indices
+        extra_args = (
+            {}
+            if (current_platform.is_mps() or current_platform.is_musa())
+            else dict(device_id=device_id)
+        )
        torch.distributed.init_process_group(
            backend=backend,
            init_method=distributed_init_method,