2P qwenimage inference, failed

[echosyy]

I tried to use two cards to infer qwen-image-edit-plus, but it failed. Could you please help to check? Thank you.

config:

GPU_memory_mode = ""
ulysses_degree = 1
ring_degree = 2
fsdp_dit = True
fsdp_text_encoder = False

compile_dit = False
enable_teacache = False

command：torchrun --nproc-per-node=2 ./examples/qwenimage/predict_t2i_edit_plus.py

error info:
0%| | 0/50 [00:00<?, ?it/s]
[rank1]: Traceback (most recent call last):
[rank1]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/./examples/qwenimage/predict_t2i_edit_plus.py", line 226, in
[rank1]: sample = pipeline(
[rank1]: ^^^^^^^^^
[rank1]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank1]: return func(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/pipeline/pipeline_qwenimage_edit_plus.py", line 864, in call
[rank1]: noise_pred = self.transformer.forward_bs(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/utils/cfg_optimization.py", line 32, in wrapper
[rank1]: result = func(self, new_x, *new_args, **new_kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/models/qwenimage_transformer2d.py", line 711, in forward_bs
[rank1]: cond_out = func(x_i, *args_i, **kwargs_i)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/models/qwenimage_transformer2d.py", line 795, in forward
[rank1]: hidden_states = self.img_in(hidden_states)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank1]: return F.linear(input, self.weight, self.bias)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: setStorage: sizes [64, 3072], strides [1, 64], storage offset 10233344, and itemsize 2 requiring a storage size of 20859904 are out of bounds for storage of size 0
[rank0]: Traceback (most recent call last):
[rank0]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/./examples/qwenimage/predict_t2i_edit_plus.py", line 226, in
[rank0]: sample = pipeline(
[rank0]: ^^^^^^^^^
[rank0]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/pipeline/pipeline_qwenimage_edit_plus.py", line 864, in call
[rank0]: noise_pred = self.transformer.forward_bs(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/utils/cfg_optimization.py", line 32, in wrapper
[rank0]: result = func(self, new_x, *new_args, **new_kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/models/qwenimage_transformer2d.py", line 711, in forward_bs
[rank0]: cond_out = func(x_i, *args_i, **kwargs_i)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/workspace/aop_lab/app_source/VideoX-Fun/videox_fun/models/qwenimage_transformer2d.py", line 795, in forward
[rank0]: hidden_states = self.img_in(hidden_states)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/admin/miniconda3/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank0]: return F.linear(input, self.weight, self.bias)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: setStorage: sizes [64, 3072], strides [1, 64], storage offset 10233344, and itemsize 2 requiring a storage size of 20859904 are out of bounds for storage of size 0

[bubbliiiing]

qwenimage不支持fsdp dit和fsdp textencoder，尚未解决。
这个pr修复了teacache的bug，#415

[bubbliiiing]

在pr中，使用如下参数应该可以正常跑，但fsdp dit和fsdp textencoder的问题还需要时间看看。
fsdp_dit = False
fsdp_text_encoder = False