Fix BLOOM crash with static cache / compiled generate

@@ -20,7 +20,7 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss from torch.nn import functional as F -from ...cache_utils import Cache, DynamicCache, StaticCache+from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer@@ -42,6 +42,22 @@ logger = logging.get_logger(__name__)  +def _get_alibi_slopes(num_heads: int, device: torch.device) -> torch.Tensor:+    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))+    base = torch.full((), 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)+    powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)+    slopes = torch.pow(base, powers)++    if closest_power_of_2 != num_heads:+        extra_base = torch.full(+            (), 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32+        )+        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)+        extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)+        slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)+    return slopes++ def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:     """     Link to paper: https://huggingface.co/papers/2108.12409 Alibi tensor is not causal as the original paper mentions, it@@ -60,23 +76,7 @@ def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torc             dtype of the output tensor     """     batch_size, seq_length = attention_mask.shape-    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))-    base = torch.full(-        (), 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32-    )-    powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)-    slopes = torch.pow(base, powers)--    if closest_power_of_2 != num_heads:-        extra_base = torch.full(-            (),-            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),-            device=attention_mask.device,-            dtype=torch.float32,-        )-        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)-        extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)-        slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)+    slopes = _get_alibi_slopes(num_heads, attention_mask.device)      # Note: alibi will added to the attention bias that will be applied to the query, key product of attention     # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)@@ -503,11 +503,20 @@ def forward(          # Compute alibi tensor: check build_alibi_tensor documentation         if attention_mask is None:-            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)+            # Without padding, alibi only depends on the key length (the whole cache length for static caches), so we+            # build it directly instead of materializing an all-ones mask+            kv_length = (+                past_key_values.get_mask_sizes(seq_length, 0)[0]+                if past_key_values is not None+                else seq_length_with_past+            )+            slopes = _get_alibi_slopes(self.num_heads, hidden_states.device)+            alibi = slopes[:, None, None] * torch.arange(kv_length, device=hidden_states.device)+            alibi = alibi.repeat(batch_size, 1, 1).to(hidden_states.dtype)         else:             attention_mask = attention_mask.to(hidden_states.device)+            alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype) -        alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)         causal_mask = create_causal_mask(             config=self.config,             inputs_embeds=inputs_embeds,@@ -571,40 +580,22 @@ def __init__(self, config: BloomConfig):     def set_output_embeddings(self, new_embeddings: torch.Tensor):         self.lm_head = new_embeddings -    def prepare_inputs_for_generation(+    def create_masks_for_generate(         self,-        input_ids,-        past_key_values=None,-        attention_mask=None,-        inputs_embeds=None,-        use_cache=True,-        is_first_iteration=False,+        config: BloomConfig,+        inputs_embeds: torch.Tensor,+        attention_mask: torch.Tensor | None,+        past_key_values: Cache | None,         **kwargs,-    ):-        # Overwritten because of the fixed-shape attention mask creation--        model_inputs = super().prepare_inputs_for_generation(-            input_ids,-            past_key_values=past_key_values,-            attention_mask=attention_mask,-            inputs_embeds=inputs_embeds,-            use_cache=use_cache,-            is_first_iteration=is_first_iteration,-            **kwargs,-        )--        # This part differs from other models because BLOOM needs a 2D mask to construct alibi tensor-        # The only difference is the usage of 2D instead of 4D mask, but the shape will be static-        if isinstance(past_key_values, StaticCache) and attention_mask is not None:-            target_length = past_key_values.get_max_length()-            batch_size, seq_length = attention_mask.shape-            diff = target_length - seq_length--            new_attn_mask = torch.zeros(batch_size, diff, device=attention_mask.device, dtype=attention_mask.dtype)-            attention_mask = torch.cat([attention_mask, new_attn_mask], dim=-1)-            model_inputs["attention_mask"] = attention_mask--        return model_inputs+    ) -> torch.Tensor:+        # Overwritten because BLOOM needs a 2D mask to build the alibi tensor, while the generic function returns a+        # 4D one. The shape stays static (the cache max length), so compiled forwards are not affected.+        if attention_mask is None or attention_mask.ndim == 4:+            return attention_mask+        target_length = past_key_values.get_max_length()+        if (padding_length := target_length - attention_mask.shape[-1]) > 0:+            attention_mask = F.pad(attention_mask, (0, padding_length))+        return attention_mask      @auto_docstring     def forward(

@@ -20,6 +20,7 @@ from transformers.testing_utils import require_torch, require_torch_accelerator, slow, torch_device  from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester+from ...generation.test_utils import assert_similar_generate_outputs from ...test_modeling_common import ids_tensor  @@ -194,6 +195,23 @@ def test_bloom_weight_initialization(self):     def test_custom_4d_attention_mask(self):         pass +    def test_generate_static_cache_without_attention_mask(self):+        # Regression test: `attention_mask=None` with a static cache used to pass a 4D mask to `build_alibi_tensor`+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()+        model = BloomForCausalLM(config).to(torch_device).eval()+        input_ids = inputs_dict["input_ids"]+        generation_kwargs = {+            "max_new_tokens": 5,+            "do_sample": False,+            "pad_token_id": None,+            "return_dict_in_generate": True,+            "output_scores": True,+        }++        dynamic_output = model.generate(input_ids, **generation_kwargs)+        static_output = model.generate(input_ids, cache_implementation="static", **generation_kwargs)+        assert_similar_generate_outputs(dynamic_output, static_output)+  @require_torch class BloomIntegrationTest(unittest.TestCase):