apply partition_spec

Author: zpcore

examples/text_to_image/train_text_to_image_xla.py

@@ -385,7 +385,7 @@ def main(args):
     server = xp.start_server(9012)
 
     num_devices = xr.global_runtime_device_count()
-    mesh = xs.get_1d_mesh('x')
+    mesh = xs.get_1d_mesh('data')
     xs.set_global_mesh(mesh)
 
     text_encoder = CLIPTextModel.from_pretrained(
@@ -521,9 +521,9 @@ def collate_fn(examples):
         device,
         input_sharding={
             "pixel_values": xs.ShardingSpec(
-                mesh, ("x", None, None, None), minibatch=True
+                mesh, ("data", None, None, None), minibatch=True
             ),
-            "input_ids": xs.ShardingSpec(mesh, ("x", None), minibatch=True),
+            "input_ids": xs.ShardingSpec(mesh, ("data", None), minibatch=True),
         },
         loader_prefetch_size=args.loader_prefetch_size,
         device_prefetch_size=args.device_prefetch_size,

src/diffusers/models/attention_processor.py

@@ -25,6 +25,7 @@
 from ..utils.torch_utils import is_torch_version, maybe_allow_in_graph
 
 if is_torch_xla_available():
+    import torch_xla.distributed.spmd as xs
     from torch_xla.experimental.custom_kernel import flash_attention
     XLA_AVAILABLE = True
 else:
@@ -2379,7 +2380,7 @@ def __call__(
                 # Apply attention mask to key
                 key = key + attention_mask
             query /= math.sqrt(query.shape[3])
-            hidden_states = flash_attention(query, key, value, causal=False)
+            hidden_states = flash_attention(query, key, value, causal=False, partition_spec=('data', None, None, None))
         else:
             hidden_states = F.scaled_dot_product_attention(
                 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False