Error when convert text_encoder to FSDP in multi gpus

hi, it's a great work!
However, I encountered an exception with script https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.1_fun/predict_i2v.py.
When generate video with multi gpus, if set fsdp_text_encoder=Truehttps://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.1_fun/predict_i2v.py#L51, I encountered an KeyError:

Traceback (most recent call last):
File "/mnt/vlab/peihao/dev_code/VideoX-Fun-main/examples/wan2.1_fun/predict_t2v.py", line 270, in
sample = pipeline(
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/mnt/vlab/peihao/dev_code/VideoX-Fun-main/videox_fun/pipeline/pipeline_wan_fun_inpaint.py", line 546, in call
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
File "/mnt/vlab/peihao/dev_code/VideoX-Fun-main/videox_fun/pipeline/pipeline_wan_fun_inpaint.py", line 280, in encode_prompt
prompt_embeds = self._get_t5_prompt_embeds(
File "/mnt/vlab/peihao/dev_code/VideoX-Fun-main/videox_fun/pipeline/pipeline_wan_fun_inpaint.py", line 223, in _get_t5_prompt_embeds
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/accelerate/hooks.py", line 175, in new_forward
output = module._old_forward(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 849, in forward
output = self._fsdp_wrapped_module(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/accelerate/hooks.py", line 170, in new_forward
args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/accelerate/hooks.py", line 341, in pre_forward
value = self.weights_map[name]
File "/mnt/vlab/peihao/miniconda3/envs/videofun2/lib/python3.10/site-packages/accelerate/utils/offload.py", line 118, in getitem
return self.dataset[f"{self.prefix}{key}"]
KeyError: '_fsdp_wrapped_module._flat_param'

But, multi-gpus inference works well, while only set fsdp_dit=Truehttps://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.1_fun/predict_i2v.py#L50, which makes me confused.

Do you have any idea about the exception?

Start with examples/wan2.1_fun/predict_i2v.py, especially the fsdp_text_encoder and fsdp_dit settings, then trace videox_fun/pipeline/pipeline_wan_fun_inpaint.py at _get_t5_prompt_embeds, encode_prompt, and the pipeline call. Reproduce multi-GPU inference with each setting and inspect the FSDP/offload interaction behind the missing _flat_param key. Done means fsdp_text_encoder no longer raises the reported KeyError and inference completes.