Fix Graph Breaks When Compiling CogView4 #10959
Conversation
Eliminate this: ``` t]V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] Recompiling function forward in /home/zeyi/repos/diffusers/src/diffusers/models/transformers/transformer_cogview4.py:374 V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] triggered by the following guard failure(s): V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] - 0/3: ___check_obj_id(L['self'].rope.freqs_h, 139976127328032) V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] - 0/2: ___check_obj_id(L['self'].rope.freqs_h, 139976107780960) V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] - 0/1: ___check_obj_id(L['self'].rope.freqs_h, 140022511848960) V0304 10:24:23.421000 3131076 torch/_dynamo/guards.py:2813] [0/4] [__recompiles] - 0/0: ___check_obj_id(L['self'].rope.freqs_h, 140024081342416) ```
| @@ -252,20 +252,18 @@ def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], th | |||
| w_inv_freq = 1.0 / (theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w)) | |||
| h_seq = torch.arange(self.rope_axes_dim[0]) | |||
| w_seq = torch.arange(self.rope_axes_dim[1]) | |||
| self.freqs_h = torch.outer(h_seq, h_inv_freq) | |||
| self.freqs_w = torch.outer(w_seq, w_inv_freq) | |||
| self.freqs_h = torch.nn.Buffer(torch.outer(h_seq, h_inv_freq)) | |||
There was a problem hiding this comment.
I think these need a persistent=False to not be part of state dict, no? Other than that, changes LGTM!
There was a problem hiding this comment.
Also would maybe prefer using register_buffer here to be consistent with some other implementations
There was a problem hiding this comment.
@a-r-r-o-w Yep, just fixed!
| @@ -252,20 +252,18 @@ def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], th | |||
| w_inv_freq = 1.0 / (theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w)) | |||
| h_seq = torch.arange(self.rope_axes_dim[0]) | |||
| w_seq = torch.arange(self.rope_axes_dim[1]) | |||
| self.freqs_h = torch.outer(h_seq, h_inv_freq) | |||
| self.freqs_w = torch.outer(w_seq, w_inv_freq) | |||
| self.freqs_h = self.register_buffer("freqs_h", torch.outer(h_seq, h_inv_freq), persistent=False) | |||
There was a problem hiding this comment.
Sorry about the back and forth but I just remembered that we did it this way so that freqs was always in float32. Using a buffer makes the tensor part of the module-modifiable parameters, so if someone were to load the model with torch_dtype=bfloat16 or do model.to(some_other_dtype), it would change the dtype of freqs. I believe that's problematic since RoPE must in fp32.
Do you think there's another way around this to avoid recompiling?
There was a problem hiding this comment.
@a-r-r-o-w I see. I am thinking about a proper workaround.
There was a problem hiding this comment.
Sorry about the back and forth but I just remembered that we did it this way so that freqs was always in
float32. Using a buffer makes the tensor part of the module-modifiable parameters, so if someone were to load the model withtorch_dtype=bfloat16or domodel.to(some_other_dtype), it would change the dtype of freqs. I believe that's problematic since RoPE must in fp32.Do you think there's another way around this to avoid recompiling?
I have made these tensors for indexing generated on the fly during the inference. And doing so seems to even make the inference a bit faster.
There was a problem hiding this comment.
. And doing so seems to even make the inference a bit faster.
this is true for both compiled and non-compiled?
There was a problem hiding this comment.
It goes against intuition for me tbh (so I will try to dive in and understand on weekend). Previously we were only computing the freqs once at initialization but now do it at each inference step, so my first thought is that is must be slower.
Since the computations are running on CPU here, including freqs-related matmul, it is understandable that it might not cause additional slowdown compared to before, since CPU is scheduling instructions much faster than GPU is processing them.
But it is followed by indexing a CPU tensor with GPU tensor, so it should introduce a CPU sync here I think. Atleast the same/similar thing was happening with our schedulers. CPU sync would drastically slow down inference but seems like it's not the case... which is interesting...
There was a problem hiding this comment.
@yiyixuxu @a-r-r-o-w I think both generating cpu tensors or cuda tensors on the fly are faster than reading from the saved one at least when compiling.🧐
|
@bot /style |
|
Style fixes have been applied. View the workflow run here. |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Eliminate this:
What does this PR do?
Fixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.