Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for fp8_e4m3fn model #80

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Add support for fp8_e4m3fn model #80

wants to merge 6 commits into from
+213 −28

Conversation

YexiongLin
Copy link

Add support for the fp8_e4m3fn model, enabling Wan2.1-480P to run on a single 4090.
Add offload VAE model, reducing occupied GPU memory.

@jinqiupeter
Copy link

Thanks for PR! I've tested 480p I2V with 1 RTX 3090 and 2 RTX 3090, both took ~35 min.

@jinqiupeter jinqiupeter mentioned this pull request Feb 27, 2025
Open
@yincangshiwei
Copy link

Gradio doesn't seem to support it yet, and Gradio's i2v-720p can't run on 48GB of video memory, so it's just OOM.

@hicaonmb
Copy link

how to modify t5.py and vae.py in order to use bf16_vae model and fp8 t5 model？there is no t5 tokenizer. json in the modelscope,where to find t5 tokenizer. json?

@YexiongLin
Copy link
Author

how to modify t5.py and vae.py in order to use bf16_vae model and fp8 t5 model？there is no t5 tokenizer. json in the modelscope,where to find t5 tokenizer. json?

I have updated the code to use fp8 t5 model. I downloaded the t5 tokenizer.json on the HuggingFace https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B/blob/main/google/umt5-xxl/tokenizer.json

@pikachurus pikachurus mentioned this pull request Mar 1, 2025
Open
@Bananamax
Copy link

Will there be fp4 support to boost Blackwell gpus?

@jffu
Copy link

jffu commented Mar 4, 2025

Thanks for your great work! I can run FP8 model on low VRAM device now.
Is multi-GPU supported? I tried enabling FSDP but encountered an error.
My command: torchrun --nproc_per_node=8 generate.py --task t2v-14B --size 720*1280 --ckpt_dir Wan2.1-T2V-14B-FP8 --dit_fsdp --t5_fsdp --ulysses_size 8 --fp8 --prompt "test" 

[2025-03-04 14:39:33,250] INFO: Creating WanT2V pipeline.
[2025-03-04 14:39:33,250] INFO: loading /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8/models_t5_umt5-xxl-enc-bf16.pth
[2025-03-04 14:40:26,251] INFO: loading /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8/Wan2.1_VAE.pth
[2025-03-04 14:40:26,560] INFO: Creating WanModel from /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8
[rank7]: Traceback (most recent call last):
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/test.py", line 366, in <module>
[rank7]:     generate(args)
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/test.py", line 319, in generate
[rank7]:     wan_t2v = wan.WanT2V(
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/wan/text2video.py", line 161, in __init__
[rank7]:     self.model = shard_fn(self.model)
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/wan/distributed/fsdp.py", line 20, in shard_model
[rank7]:     model = FSDP(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 483, in __init__
[rank7]:     _auto_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 101, in _auto_wrap
[rank7]:     _recursive_wrap(**recursive_wrap_kwargs, **root_kwargs)  # type: ignore[arg-type]
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 545, in _recursive_wrap
[rank7]:     wrapped_child, num_wrapped_params = _recursive_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 545, in _recursive_wrap
[rank7]:     wrapped_child, num_wrapped_params = _recursive_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 563, in _recursive_wrap
[rank7]:     return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 492, in _wrap
[rank7]:     return wrapper_cls(module, **kwargs)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 509, in __init__
[rank7]:     _init_param_handle_from_module(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 636, in _init_param_handle_from_module
[rank7]:     _init_param_handle_from_params(state, managed_params, fully_sharded_module)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 648, in _init_param_handle_from_params
[rank7]:     handle = FlatParamHandle(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 602, in __init__
[rank7]:     self._init_flat_param_and_metadata(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 652, in _init_flat_param_and_metadata
[rank7]:     ) = self._validate_tensors_to_flatten(params)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 796, in _validate_tensors_to_flatten
[rank7]:     raise ValueError(
[rank7]: ValueError: Must flatten tensors with uniform dtype but got torch.bfloat16 and torch.float8_e4m3fn

@YexiongLin
Copy link
Author

Thanks for your great work! I can run FP8 model on low VRAM device now. Is multi-GPU supported? I tried enabling FSDP but encountered an error. My command: torchrun --nproc_per_node=8 generate.py --task t2v-14B --size 720*1280 --ckpt_dir Wan2.1-T2V-14B-FP8 --dit_fsdp --t5_fsdp --ulysses_size 8 --fp8 --prompt "test" 

[2025-03-04 14:39:33,250] INFO: Creating WanT2V pipeline.
[2025-03-04 14:39:33,250] INFO: loading /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8/models_t5_umt5-xxl-enc-bf16.pth
[2025-03-04 14:40:26,251] INFO: loading /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8/Wan2.1_VAE.pth
[2025-03-04 14:40:26,560] INFO: Creating WanModel from /home/admin/workspace/aop_lab/app_data/Wan2.1-T2V-14B-FP8
[rank7]: Traceback (most recent call last):
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/test.py", line 366, in <module>
[rank7]:     generate(args)
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/test.py", line 319, in generate
[rank7]:     wan_t2v = wan.WanT2V(
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/wan/text2video.py", line 161, in __init__
[rank7]:     self.model = shard_fn(self.model)
[rank7]:   File "/home/admin/workspace/aop_lab/app_source/Wan2.1-FP8/wan/distributed/fsdp.py", line 20, in shard_model
[rank7]:     model = FSDP(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 483, in __init__
[rank7]:     _auto_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 101, in _auto_wrap
[rank7]:     _recursive_wrap(**recursive_wrap_kwargs, **root_kwargs)  # type: ignore[arg-type]
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 545, in _recursive_wrap
[rank7]:     wrapped_child, num_wrapped_params = _recursive_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 545, in _recursive_wrap
[rank7]:     wrapped_child, num_wrapped_params = _recursive_wrap(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 563, in _recursive_wrap
[rank7]:     return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/wrap.py", line 492, in _wrap
[rank7]:     return wrapper_cls(module, **kwargs)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 509, in __init__
[rank7]:     _init_param_handle_from_module(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 636, in _init_param_handle_from_module
[rank7]:     _init_param_handle_from_params(state, managed_params, fully_sharded_module)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_init_utils.py", line 648, in _init_param_handle_from_params
[rank7]:     handle = FlatParamHandle(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 602, in __init__
[rank7]:     self._init_flat_param_and_metadata(
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 652, in _init_flat_param_and_metadata
[rank7]:     ) = self._validate_tensors_to_flatten(params)
[rank7]:   File "/home/admin/miniconda3/lib/python3.10/site-packages/torch/distributed/fsdp/_flat_param.py", line 796, in _validate_tensors_to_flatten
[rank7]:     raise ValueError(
[rank7]: ValueError: Must flatten tensors with uniform dtype but got torch.bfloat16 and torch.float8_e4m3fn

Due to the parameters of model using mixed precision, FSDP is not supported currently

@pikachurus pikachurus mentioned this pull request Mar 7, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@YexiongLin @jinqiupeter @yincangshiwei @hicaonmb @Bananamax @jffu