[diffusion] multi-platform: use current_platform.device_type to replace hard-coded cuda device (#15232)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
2025-12-16 22:17:33 +08:00
parent ecb401ed42
commit 9f4ed93dd8
16 changed files with 62 additions and 23 deletions
@@ -448,7 +448,7 @@ class TextEncoderLoader(ComponentLoader):
                    )
                else:
                    mesh = init_device_mesh(
-                        "cuda",
+                        current_platform.device_type,
                        mesh_shape=(1, dist.get_world_size()),
                        mesh_dim_names=("offload", "replicate"),
                    )
@@ -116,7 +116,7 @@ def maybe_load_fsdp_model(
            hsdp_shard_dim = 1

        device_mesh = init_device_mesh(
-            "cuda",
+            current_platform.device_type,
            # (Replicate(), Shard(dim=0))
            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
            mesh_dim_names=("replicate", "shard"),