Disable CFG cache when batch size is indivisible by data shards to prevent ulysses attention crashes

Author: Perseus14

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -256,6 +256,19 @@ def run_inference_2_2(
   do_classifier_free_guidance = guidance_scale_low > 1.0 or guidance_scale_high > 1.0
   bsz = latents.shape[0]
 
+  data_shards = 1
+  try:
+    if hasattr(latents, "sharding") and hasattr(latents.sharding, "mesh"):
+      data_shards = latents.sharding.mesh.shape["data"] * latents.sharding.mesh.shape.get("fsdp", 1)
+  except Exception:
+    pass
+
+  if use_cfg_cache and do_classifier_free_guidance and bsz % data_shards != 0:
+    from maxdiffusion.max_logging import MaxLogging
+    max_logging = MaxLogging()
+    max_logging.log(f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1.")
+    use_cfg_cache = False
+
   prompt_embeds_combined = (
       jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0) if do_classifier_free_guidance else prompt_embeds
   )

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -371,6 +371,19 @@ def run_inference_2_2_i2v(
   do_classifier_free_guidance = guidance_scale_low > 1.0 or guidance_scale_high > 1.0
   bsz = latents.shape[0]
 
+  data_shards = 1
+  try:
+    if hasattr(latents, "sharding") and hasattr(latents.sharding, "mesh"):
+      data_shards = latents.sharding.mesh.shape["data"] * latents.sharding.mesh.shape.get("fsdp", 1)
+  except Exception:
+    pass
+
+  if use_cfg_cache and do_classifier_free_guidance and bsz % data_shards != 0:
+    from maxdiffusion.max_logging import MaxLogging
+    max_logging = MaxLogging()
+    max_logging.log(f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1.")
+    use_cfg_cache = False
+
   prompt_embeds_combined = (
       jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0) if do_classifier_free_guidance else prompt_embeds
   )