wan2.2 multi-gpu inference issue

[LecterF]

VAE encoding: 100%|██████████| 1/1 [00:00<00:00, 1.12it/s] VAE encoding: 100%|██████████| 1/1 [00:00<00:00, 1.16it/s] VAE encoding: 100%|██████████| 1/1 [00:00<00:00, 1.15it/s] VAE encoding: 100%|██████████| 1/1 [00:00<00:00, 1.11it/s] 0%| | 0/50 [00:00<?, ?it/s] [rank1]: Traceback (most recent call last): [rank1]: File "/data2/users/fengxiaoyi/video_gen/Index-anisora/models/test_2_2.py", line 100, in <module> [rank1]: video = pipe( [rank1]: ^^^^^ [rank1]: File "/data2/users/fengxiaoyi/conda_env/video/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context [rank1]: return func(*args, **kwargs) [rank1]: ^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/video_gen/Index-anisora/DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py", line 449, in __call__ [rank1]: noise_pred_posi = self.model_fn(**models, **inputs_shared, **inputs_posi, timestep=timestep) [rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/video_gen/Index-anisora/DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py", line 1101, in model_fn_wan_video [rank1]: x = block(x, context, t_mod, freqs) [rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/conda_env/video/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl [rank1]: return self._call_impl(*args, **kwargs) [rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/conda_env/video/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl [rank1]: return forward_call(*args, **kwargs) [rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/video_gen/Index-anisora/DiffSynth-Studio/diffsynth/models/wan_video_dit.py", line 225, in forward [rank1]: input_x = modulate(self.norm1(x), shift_msa, scale_msa) [rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [rank1]: File "/data2/users/fengxiaoyi/video_gen/Index-anisora/DiffSynth-Studio/diffsynth/models/wan_video_dit.py", line 65, in modulate [rank1]: return (x * (1 + scale) + shift) [rank1]: ~~^~~~~~~~~~~~~ [rank1]: RuntimeError: The size of tensor a (878) must match the size of tensor b (3510) at non-singleton dimension 1 0%| | 0/50 [00:00<?, ?it/s]

用的wan2.2-5B进行推理，单卡可以跑，但是多卡会出现这样的错误

[hrdxwandg]

I2V-A14B 多卡推理同样报错，单卡没问题

[rank4]: Traceback (most recent call last):
[rank4]:   File "xxx/Wan2.2-I2V-A14B_inference.py", line 33, in <module>
[rank4]:     video = pipe(
[rank4]:             ^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank4]:     return func(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/diffsynth/pipelines/wan_video_new.py", line 449, in __call__
[rank4]:     noise_pred_posi = self.model_fn(**models, **inputs_shared, **inputs_posi, timestep=timestep)
[rank4]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/diffsynth/pipelines/wan_video_new.py", line 1107, in model_fn_wan_video
[rank4]:     x = block(x, context, t_mod, freqs)
[rank4]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank4]:     return self._call_impl(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank4]:     return forward_call(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/diffsynth/models/wan_video_dit.py", line 226, in forward
[rank4]:     x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
[rank4]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank4]:     return self._call_impl(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank4]:     return forward_call(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/diffsynth/distributed/xdit_context_parallel.py", line 121, in usp_attn_forward
[rank4]:     x = xFuserLongContextAttention()(
[rank4]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank4]:     return self._call_impl(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank4]:     return forward_call(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
[rank4]:     return fn(*args, **kwargs)
[rank4]:            ^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/xfuser/core/long_ctx_attention/hybrid/attn_layer.py", line 176, in forward
[rank4]:     query_layer = SeqAllToAll4D.apply(
[rank4]:                   ^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/torch/autograd/function.py", line 575, in apply
[rank4]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/yunchang/comm/all_to_all.py", line 122, in forward
[rank4]:     return all_to_all_4D(input, scatter_idx, gather_idx, group=group, use_sync=use_sync)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "xxx/envs/diffsynth/lib/python3.11/site-packages/yunchang/comm/all_to_all.py", line 46, in all_to_all_4D
[rank4]:     input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs)
[rank4]: RuntimeError: shape '[1, 5460, 6, 6, 128]' is invalid for input of size 27955200

[mi804]

The cause of this issue is that the sequence length of attention cannot be evenly divided by the number of GPUs. This issue has been fixed in the latest source code. Please update the diffsynth-studio source code. Of course, we still recommend using a number of GPUs that can be evenly divided, such as 4.

[hrdxwandg]

The cause of this issue is that the sequence length of attention cannot be evenly divided by the number of GPUs. This issue has been fixed in the latest source code. Please update the diffsynth-studio source code. Of course, we still recommend using a number of GPUs that can be evenly divided, such as 4.

ok, thanks

[LecterF]

The cause of this issue is that the sequence length of attention cannot be evenly divided by the number of GPUs. This issue has been fixed in the latest source code. Please update the diffsynth-studio source code. Of course, we still recommend using a number of GPUs that can be evenly divided, such as 4.

Thanks！

[JianxGao]

It seems that the current code still encounters a shape error: x is split across multiple GPUs after use_unified_sequence_parallel, but t_mod is computed based on the original shape of x. Simply applying a chunk operation to t_mod after splitting x does not resolve the issue.

[JianxGao]

The following code seems to fix the issue, but I'm not sure if it's the right way to handle it.

t_mod = torch.chunk(t_mod, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
t = torch.chunk(t, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]

DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py

Lines 1075 to 1080 in 32cf5d3

	if use_unified_sequence_parallel:
	if dist.is_initialized() and dist.get_world_size() > 1:
	chunks = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)
	pad_shape = chunks[0].shape[1] - chunks[-1].shape[1]
	chunks = [torch.nn.functional.pad(chunk, (0, 0, 0, chunks[0].shape[1]-chunk.shape[1]), value=0) for chunk in chunks]
	x = chunks[get_sequence_parallel_rank()]

[mi804]

@JianxGao
The USP is split along the sequence dimension, while t_mod is independent of the sequence dimension, so t_mod does not need to be split in USP. The issue you're encountering is most likely due to underlying limitations in xdit. Before performing attention, xdit gathers all tensors and then splits them along the num_heads dimension. Therefore, the number of your parallel GPUs must evenly divide the num_heads (for Wan2.2 5B, num_heads=24). Otherwise, an error will occur.

Of course, feel free to share your parallel inference code and the number of GPUs you're using—we'll follow up with further analysis.

[JianxGao]

Hi @mi804, thanks for the quick reply!

I'm using Wan2.2-TI2V-5B to generate a video at a resolution of 704×1280. During generation, both dit.separated_timestep and fuse_vae_embedding_in_latents are set to True.

DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py

Lines 1017 to 1024 in 32cf5d3

	# Timestep
	if dit.seperated_timestep and fuse_vae_embedding_in_latents:
	timestep = torch.concat([
	torch.zeros((1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device),
	torch.ones((latents.shape[2] - 1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device) * timestep
	]).flatten()
	t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep).unsqueeze(0))
	t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim))

As a result, the shape of timesep is [18400], so the shape of t_mod becomes [1, 18400, 6, 3072].
Hence, t_mod does have a spatial dimension.

[mi804]

Sorry that I have ignored the special case when using T2V mode of WAN2.2 5B, I will fix this soon

[JianxGao]

Thanks very much. However, I infer the model with I2V mode and haven’t tested T2V mode yet. It's possible that T2V may exhibit a similar issue.

[mi804]

@JianxGao This should have been fixed in #763. You may try and let us know if there are other issues 😄.