NaN miscompile: async-collective-fusion + overlap-compute-collective corrupt Wan2.2-I2V ulysses+splash attention on v6e (blocks ~5% DiT speedup)

[ThomasNing]

Description

On TPU v6e-16 (multi-host), combining two XLA:TPU collective-overlap flag groups silently corrupts outputs to NaN for a sequence-parallel (Ulysses) + splash-attention workload (MaxDiffusion Wan2.2-I2V-A14B DiT). Each group alone is numerically correct; only together do they NaN — while delivering a real ~5% speedup. The corruption appears tied to overlapping the Ulysses all-to-all/all-reduce collectives around the splash-attention Mosaic custom-call.

Environment

• TPU v6e-16 (4 hosts × 4 chips), multi-host via jax.distributed
• jax 0.10.2, flax 0.12.7, recent libtpu; bfloat16
• MaxDiffusion Wan2.2-I2V-A14B (two-expert DiT), attention=ulysses_custom (splash/flash), ici_context=4 × ici_tensor=4

Flags that trigger it (LIBTPU_INIT_ARGS)

--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_reduce=true

Symptom

The DiT denoise output becomes NaN (and diverges across processes); caught downstream by a jax.device_put → multihost_utils.assert_equal on a replicated sharding:

AssertionError: ... not the same on each process. Expected: [[[[[nan nan nan ...

Bisection (warm DiT denoise; "correct" = downstream VAE decode succeeds)

flags	warm DiT	correct?
baseline (none)	269.7 s	✅
async-collective-fusion group only	261.6 s	✅
overlap-compute-collective group only	262.1 s	✅
megacore-fusion only	260.9 s	✅
fusion + megacore	260.6 s	✅
overlap + megacore	262.0 s	✅
fusion + overlap	247.9 s	❌ NaN
enable_async_collective_fusion + overlap_compute_collective_tc (primaries only)	260.7 s	✅

So it's specifically async-collective-fusion × overlap-compute-collective together, and requires the aggressive sub-flags (fuse_all_reduce / multiple_steps / async_all_reduce) that overlap collectives across steps. The ~5% speedup appears only in the NaN-producing config.

What we ruled out (no minimal repro yet)

A standalone multi-host program on the same 4×4 v6e mesh — a loop of matmul → psum(all-reduce) → all_to_all → all_to_all → tanh — does not reproduce, in f32 or bf16 (output bit-identical with vs without the flags). So this is not generic collective+overlap; it needs the workload's specific collective structure — most likely the interaction with the splash-attention Mosaic custom-call in the Ulysses attention path (which a pure-XLA matmul graph lacks).

Ask

Is this a known issue with async-collective-fusion + overlap-compute-collective around Mosaic/splash-attention custom-calls on v6e? A self-contained minimal repro would need to include the splash-kernel + Ulysses all-to-all pattern; happy to provide HLO dumps or test candidate fixes/flags with guidance.

---

Filing here (rather than jax-ml/jax) since this is MaxDiffusion's ulysses_custom + splash-attention path on v6e, and the combination delivers a real ~5% DiT speedup that's currently unusable. Questions: is this a known incompatibility for the Wan2.2-I2V Ulysses path? Should MaxDiffusion guard/validate these flags, or is there a MaxDiffusion-side scheduling fix — and can this be escalated to the XLA:TPU/libtpu team with the model repro?

[csgoogle]

Hi @ThomasNing , I think this might be happening due to allreduce. Also, out of curiosity, what’s the reasoning behind using TP here? For similar workloads, we've found that combining FSDP + CP + DP usually delivers the best results.

[ThomasNing]

Thanks for the pointer! all-reduce is also our suspicion, which makes sense, given that the failure requires the fuse_all_reduce / async_all_reduce options in addition to the overlap pass.

We have tested with several TP + CP + DP configurations. For this workload and model shape, scaling TP seems to have given us the lowest peak per-device memory usage, which allows us to use a meaningfully larger microbatch, resulting in higher end-to-end throughput. That said, we are very open to trying any recommended FSDP/CP/DP/TP configuration recipe that you believe would maximize throughput further on this workload.

[csgoogle]

Hi @ThomasNing , Can you also provide the batch size and confirm whether you’re using CFG cache? If possible, it would also be helpful if you could share the full command you’re running.

[ThomasNing]

@csgoogle
Global batch size 1 on v6e-16 (0.0625 × 16 chips), 1 video per prompt. Note that with CFG enabled the DiT itself runs a doubled batch of 2 internally (cond ‖ uncond concatenated along axis 0), so the Ulysses all-to-all + splash-attention custom-call sees batch 2 per denoise step.

CFG cache: No — use_cfg_cache=False. Classifier-free guidance itself is on, though: do_classifier_free_guidance=True via guidance_scale_low=3.0, guidance_scale_high=4.0. (I separately confirmed use_cfg_cache=True is a no-op on this i2v-2.2 path — bit-identical output and no speedup — in case that's relevant to the collective structure.)

[csgoogle]

Thanks, for bs=1 we observed dp=2 and cp=8 is the best config when cfg is enabled. You can see recipe here (it's for ironwood but most of the configuration is same)

https://github.com/AI-Hypercomputer/tpu-recipes/blob/main/inference/ironwood/MaxDiffusion/Wan2.x/Wan2.2-T2V/run_recipe.sh

[ThomasNing]

Thank you for pointing that out. I will take a look and do the experiment on my end. When delivering the service, to balance between throughput and latency, we choose the batch size = 4. In that case, dp=4 and cp=4 win.

[ThomasNing]

Thanks! Answers first, then a parallelization follow-up since it ties back to the NaN.

Batch size: global batch 1 (per_device_batch_size=0.0625 × 16 chips), 1 video/prompt. With CFG on, the DiT runs a doubled batch of 2 internally (concat([cond, uncond])), so the Ulysses all-to-all + splash custom-call sees batch 2/step.

CFG cache: no (use_cfg_cache=False); CFG itself is on (guidance_scale_low=3.0, guidance_scale_high=4.0). (Separately confirmed use_cfg_cache=True is a no-op on this i2v-2.2 path — bit-identical output.)

Config: Wan2.2-I2V-A14B from base_wan_i2v_27b.yml, 1088×1920, 81 frames, bf16, attention=ulysses_custom, mesh DP1 × CP4 × TP4 (ici_context_parallelism=4 ici_tensor_parallelism=4), v6e-16 multi-host. (Driven via a staged DiT-only runner; happy to share it / HLO dumps.)

On the dp=2/cp=8 tip — it's right, and TP is the key. I swept the parallelization at bs=4 (CFG off, 4 videos, same 16 chips):

Mesh	warm denoise (4 videos)	per-video
DP1 × CP4 × TP4 (what I had)	887 s	222 s
DP4 × CP4 × TP1	226 s	57 s (3.9×)
DP2 × CP8 × TP1	237 s	59 s

Dropping TP and scaling DP to the batch is ~3.9× more economical — TP=4 gave each video 16 chips for the same latency CP=4 alone delivers. Exactly the "drop TP, scale DP" guidance.

And this looks like the NaN fix too. My NaN trigger included fuse_all_reduce + xla_enable_async_all_reduce, and the all-reduces in my mesh came from TP=4. The Ironwood recipe runs dp×cp with tp=1 (no tensor-parallel all-reduce) and omits exactly those two all-reduce-fusion flags while keeping the overlap primaries. So moving to dp×cp (tp=1) should both win on throughput and sidestep the NaN — the corrupting interaction (TP all-reduce × overlap around the splash custom-call) doesn't exist there. Happy to re-run the bisection on the tp=1 mesh to confirm.

[csgoogle]

Yes, that's exactly right. We should only encounter all-reduce collectives when Tensor Parallelism (TP) is enabled. For a hybrid FSDP + CP configuration, the expected behavior is to see all-gather operations for the weights and all-to-all operations for the Ulysses sequence attention.