aigc-apps / aigc-apps/VideoX-Fun

How to use torch.compile() on transformer

Open
#45 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
2.3k
Forks
188
Avg merge
2d 22h
Merged PRs (30d)
3

Description

`# Text Encoder
text_encoder = T5EncoderModel.from_pretrained(model_path, subfolder="text_encoder", torch_dtype=weight_dtype).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 ?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with cog-fun-i2v.py at pipeline initialization and reproduce the failure using the shown quantize_ and torch.compile sequence on the transformer. Trace the interaction through PyTorch's Module._apply and torch.utils.swap_tensors, then compare initialization with and without compilation. Done means identifying a reproducible compatible setup or documenting the limitation and required change.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.