Skip to content

Fix Graph Breaks When Compiling CogView4 #10959

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

Merged
merged 5 commits into from
Mar 7, 2025
Merged

Fix Graph Breaks When Compiling CogView4 #10959

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7fc465f
Fix Graph Breaks When Compiling CogView4
chengzeyi Mar 4, 2025
9080bd6
Update transformer_cogview4.py
chengzeyi Mar 4, 2025
74b591a
fix cogview4 rotary pos embed
chengzeyi Mar 6, 2025
7d18503
Apply style fixes
github-actions[bot] Mar 6, 2025
732356a
Merge branch 'main' into patch-1
yiyixuxu Mar 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions src/diffusers/models/transformers/transformer_cogview4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@a-r-r-o-w I see. I am thinking about a proper workaround.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

interesting! thanks!

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

. And doing so seems to even make the inference a bit faster.

this is true for both compiled and non-compiled?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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...

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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.🧐

self.freqs_w = self.register_buffer("freqs_h", torch.outer(w_seq, w_inv_freq), persistent=False)

def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
batch_size, num_channels, height, width = hidden_states.shape
height, width = height // self.patch_size, width // self.patch_size

h_idx = torch.arange(height)
w_idx = torch.arange(width)
h_idx = torch.arange(height, device=self.freqs_h.device)
w_idx = torch.arange(width, device=self.freqs_w.device)
inner_h_idx = h_idx * self.rope_axes_dim[0] // height
inner_w_idx = w_idx * self.rope_axes_dim[1] // width

self.freqs_h = self.freqs_h.to(hidden_states.device)
self.freqs_w = self.freqs_w.to(hidden_states.device)
freqs_h = self.freqs_h[inner_h_idx]
freqs_w = self.freqs_w[inner_w_idx]

Expand Down