diff --git a/src/diffusers/pipelines/kolors/text_encoder.py b/src/diffusers/pipelines/kolors/text_encoder.py index 434f4fed6fbb..9eeeda20b5f6 100644 --- a/src/diffusers/pipelines/kolors/text_encoder.py +++ b/src/diffusers/pipelines/kolors/text_encoder.py @@ -130,101 +130,24 @@ def __init__(self, config: ChatGLMConfig, layer_number): self.attention_dropout = torch.nn.Dropout(config.attention_dropout) def forward(self, query_layer, key_layer, value_layer, attention_mask): - pytorch_major_version = int(torch.__version__.split(".")[0]) - if pytorch_major_version >= 2: - query_layer, key_layer, value_layer = [ - k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer] - ] - if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: - context_layer = torch.nn.functional.scaled_dot_product_attention( - query_layer, key_layer, value_layer, is_causal=True - ) - else: - if attention_mask is not None: - attention_mask = ~attention_mask - context_layer = torch.nn.functional.scaled_dot_product_attention( - query_layer, key_layer, value_layer, attention_mask - ) - context_layer = context_layer.permute(2, 0, 1, 3) - new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) - context_layer = context_layer.reshape(*new_context_layer_shape) - else: - # Raw attention scores - - # [b, np, sq, sk] - output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) - - # [sq, b, np, hn] -> [sq, b * np, hn] - query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) - # [sk, b, np, hn] -> [sk, b * np, hn] - key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) - - # preallocting input tensor: [b * np, sq, sk] - matmul_input_buffer = torch.empty( - output_size[0] * output_size[1], - output_size[2], - output_size[3], - dtype=query_layer.dtype, - device=query_layer.device, - ) - - # Raw attention scores. [b * np, sq, sk] - matmul_result = torch.baddbmm( - matmul_input_buffer, - query_layer.transpose(0, 1), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] - beta=0.0, - alpha=(1.0 / self.norm_factor), + # diffusers requires torch >= 2.6, so scaled_dot_product_attention is always + # available; the pre-2.0 manual attention path this module originally carried + # was unreachable (and used an MPS-unsafe baddbmm idiom, see + # https://github.com/huggingface/diffusers/issues/14624). + query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] + if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: + context_layer = torch.nn.functional.scaled_dot_product_attention( + query_layer, key_layer, value_layer, is_causal=True ) - - # change view to [b, np, sq, sk] - attention_scores = matmul_result.view(*output_size) - - # =========================== - # Attention probs and dropout - # =========================== - - # attention scores and attention mask [b, np, sq, sk] - if self.attention_softmax_in_fp32: - attention_scores = attention_scores.float() - if self.coeff is not None: - attention_scores = attention_scores * self.coeff - if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]: - attention_mask = torch.ones( - output_size[0], 1, output_size[2], output_size[3], device=attention_scores.device, dtype=torch.bool - ) - attention_mask.tril_() - attention_mask = ~attention_mask + else: if attention_mask is not None: - attention_scores = attention_scores.masked_fill(attention_mask, float("-inf")) - attention_probs = F.softmax(attention_scores, dim=-1) - attention_probs = attention_probs.type_as(value_layer) - - # This is actually dropping out entire tokens to attend to, which might - # seem a bit unusual, but is taken from the original Transformer paper. - attention_probs = self.attention_dropout(attention_probs) - # ========================= - # Context layer. [sq, b, hp] - # ========================= - - # value_layer -> context layer. - # [sk, b, np, hn] --> [b, np, sq, hn] - - # context layer shape: [b, np, sq, hn] - output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3)) - # change view [sk, b * np, hn] - value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) - # change view [b * np, sq, sk] - attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) - # matmul: [b * np, sq, hn] - context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) - # change view [b, np, sq, hn] - context_layer = context_layer.view(*output_size) - # [b, np, sq, hn] --> [sq, b, np, hn] - context_layer = context_layer.permute(2, 0, 1, 3).contiguous() - # [sq, b, np, hn] --> [sq, b, hp] - new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) - context_layer = context_layer.view(*new_context_layer_shape) + attention_mask = ~attention_mask + context_layer = torch.nn.functional.scaled_dot_product_attention( + query_layer, key_layer, value_layer, attention_mask + ) + context_layer = context_layer.permute(2, 0, 1, 3) + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.reshape(*new_context_layer_shape) return context_layer diff --git a/tests/pipelines/kolors/test_kolors_text_encoder.py b/tests/pipelines/kolors/test_kolors_text_encoder.py new file mode 100644 index 000000000000..f08909590900 --- /dev/null +++ b/tests/pipelines/kolors/test_kolors_text_encoder.py @@ -0,0 +1,74 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.pipelines.kolors.text_encoder import ChatGLMConfig, CoreAttention + +from ...testing_utils import torch_device + + +class TestKolorsCoreAttention: + """ + `CoreAttention` computes attention through `scaled_dot_product_attention`; the module + used to carry a second, manual attention path gated on torch < 2 that was unreachable + on any supported torch (https://github.com/huggingface/diffusers/issues/14624). These + tests pin the behaviour of the remaining path on every backend so its removal, and any + future rework, stay observable. + """ + + def get_attention(self): + config = ChatGLMConfig( + hidden_size=256, + num_attention_heads=8, + kv_channels=32, + multi_query_attention=False, + attention_softmax_in_fp32=True, + ) + return CoreAttention(config, layer_number=1) + + def get_inputs(self, device, dtype=torch.float32, seq_len=64, batch=2, heads=8, head_dim=32): + torch.manual_seed(0) + shape = (seq_len, batch, heads, head_dim) + return tuple(torch.randn(shape, dtype=dtype).to(device) for _ in range(3)) + + @pytest.mark.parametrize("masked", [False, True]) + def test_output_matches_cpu_reference(self, masked): + # The device path must stay numerically equivalent to the CPU path, for both the + # causal (mask=None) branch and the explicit-mask branch of forward. + attention = self.get_attention() + query, key, value = self.get_inputs("cpu") + seq_len, batch = query.shape[0], query.shape[1] + + if masked: + torch.manual_seed(1) + # ChatGLM convention: True marks positions that must NOT be attended to. + attention_mask = torch.rand(batch, 1, seq_len, seq_len) < 0.25 + attention_mask[..., 0] = False # keep at least one visible key per query row + else: + attention_mask = None + + with torch.no_grad(): + expected = attention(query, key, value, attention_mask) + actual = attention( + query.to(torch_device), + key.to(torch_device), + value.to(torch_device), + attention_mask.to(torch_device) if attention_mask is not None else None, + ) + + assert torch.isfinite(expected).all() + torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4)