aigc-apps / CogVideoX-Fun

📹 A more flexible CogVideoX that can generate videos at any resolution and creates videos from images.
Apache License 2.0
486 stars 32 forks source link

How to use torch.compile() on transformer #45

Open kursatdinc opened 1 month ago

kursatdinc commented 1 month ago

`# Text Encoder text_encoder = T5EncoderModel.from_pretrained(model_path, subfolder="text_encoder", torch_dtype=weightdtype).to(device) quantize(text_encoder, quantization())

# Transformer
transformer = CogVideoXTransformer3DModel.from_pretrained_2d(model_path, subfolder="transformer").to(weight_dtype).to(device)
quantize_(transformer, quantization())

transformer = torch.compile(
    transformer,
    mode = 'reduce-overhead',
    fullgraph=True,
    dynamic=True
    )`

When I use torch.compile initialize pipeline, throw an Couldn't swap Linear.weight error.

`---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File ~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:945, in Module._apply(self, fn, recurse)
    [942](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:942)     param_applied = torch.nn.Parameter(
    [943](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:943)         param_applied, requires_grad=param.requires_grad
    [944](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:944)     )
--> [945](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:945)     torch.utils.swap_tensors(param, param_applied)
    [946](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:946) except Exception as e:

File ~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/utils/__init__.py:51, in swap_tensors(t1, t2)
     [50](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/utils/__init__.py:50) if weakref.getweakrefs(t1):
---> [51](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/utils/__init__.py:51)     raise RuntimeError("Cannot swap t1 because it has weakref associated with it")
     [52](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/utils/__init__.py:52) if weakref.getweakrefs(t2):

RuntimeError: Cannot swap t1 because it has weakref associated with it

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
File /home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:2
      [1](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:1) start_time = time.time()
----> [2](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:2) generated_video = generate_video(
      [3](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:3)     pipeline = initialized_pipeline,
      [4](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:4)     vae = initialized_vae,
      [5](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/video_model/CogVideoX-Fun/cog-fun-i2v.py:5)     lora_path = lora_path,
...
    [951](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:951)         ) from e
    [952](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:952)     out_param = param
    [953](https://vscode-remote+ssh-002dremote-002ba100-002d40gb-002dspot.vscode-resource.vscode-cdn.net/home/kursat_dinc/desktop/~/anaconda3/envs/cogvideo/lib/python3.10/site-packages/torch/nn/modules/module.py:953) elif p_should_use_set_data:

RuntimeError: _apply(): Couldn't swap Linear.weight`

How can ı solve this problem ?

kursatdinc commented 1 month ago

And also, quantize not work correctly, there is no inference time improvement when I use quantize.