From 4f68eb7ddf27cfaf8606ac0de419e6ed3e65f60c Mon Sep 17 00:00:00 2001 From: Artur Kloniecki Date: Fri, 19 Sep 2025 12:43:20 +0300 Subject: [PATCH 1/2] Backport GPTBigCodeAttention from v4.53.0 to avoid major refactoring. Signed-off-by: Artur Kloniecki --- .../models/gpt_bigcode/modeling_gpt_bigcode.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py index 150aef1bae..5502fc668f 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -60,6 +60,20 @@ def __init__(self, config, is_cross_attention=False, layer_idx=None): self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA is not None else None self.block_size = 4096 + # Starting with v4.54 self.attn_dropout in GPTBigCodeAttention constructor is only assigned config.attn_pdrop value + # and not a Dropout operator, due to overall refactor of the class. Here we overwrite the value back with correct + # OP, to make use of the previous flow without changes. + self.attn_dropout = torch.nn.Dropout(config.attn_pdrop) + + def _get_mask_value(self, device, dtype): + """ + This method has been copied from GPT_BigCodeAttention._get_mask_value: https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py + In further releases whole class has been rewritten and is no longer compatible with our working solution. + """ + # torch.where expects a tensor. We use a cache to avoid recreating it every time. + if self.mask_value is None or self.mask_value.dtype != dtype or self.mask_value.device != device: + self.mask_value = torch.full([], torch.finfo(dtype).min, dtype=dtype, device=device) + return self.mask_value def _attn(self, query, key, value, attention_mask=None, head_mask=None): """ From 44400cd4ab263d2aac9f659e4bf076f02bbdbb5d Mon Sep 17 00:00:00 2001 From: Artur Kloniecki Date: Fri, 19 Sep 2025 19:45:55 +0300 Subject: [PATCH 2/2] Address PR comment. Signed-off-by: Artur Kloniecki --- .../models/gpt_bigcode/modeling_gpt_bigcode.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py index 5502fc668f..e9063530d4 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -60,10 +60,6 @@ def __init__(self, config, is_cross_attention=False, layer_idx=None): self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA is not None else None self.block_size = 4096 - # Starting with v4.54 self.attn_dropout in GPTBigCodeAttention constructor is only assigned config.attn_pdrop value - # and not a Dropout operator, due to overall refactor of the class. Here we overwrite the value back with correct - # OP, to make use of the previous flow without changes. - self.attn_dropout = torch.nn.Dropout(config.attn_pdrop) def _get_mask_value(self, device, dtype): """ @@ -81,6 +77,8 @@ def _attn(self, query, key, value, attention_mask=None, head_mask=None): Copied from GPTBigCodeAttention._attn: https://github.com/huggingface/transformers/blob/v4.40-release/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py The only differences are: - in self._attn, use torch.matmul instead of torch.baddbmm when the device used for query is not cpu + - starting with v4.54, self.attn_dropout is a scalar value, assigned from config in constructor, and not a Dropout operator; + use torch.nn.functional.dropout() in it's place to restore previous functionality. """ dtype = query.dtype softmax_dtype = torch.float32 if self.attention_softmax_in_fp32 else dtype @@ -144,7 +142,7 @@ def _attn(self, query, key, value, attention_mask=None, head_mask=None): attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) - attn_weights = self.attn_dropout(attn_weights) + attn_weights = torch.nn.functional.dropout(attn_weights, p=self.attn_dropout) # Mask heads if we want to if head_mask is not None: