fix AutoTP in deepspeed could not work for bloom (huggingface#8)

Signed-off-by: Wang, Yi A <[email protected]>
sywangyi · Mar 16, 2023 · bc46b47 · bc46b47
1 parent 1c4a9ac
commit bc46b47
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/models/bloom/modeling_bloom.py b/src/transformers/models/bloom/modeling_bloom.py
@@ -638,6 +638,8 @@ def __init__(self, config: BloomConfig):
 
         self.gradient_checkpointing = False
 
+        self.build_alibi_tensor = build_alibi_tensor
+
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -750,7 +752,7 @@ def forward(
         else:
             attention_mask = attention_mask.to(hidden_states.device)
 
-        alibi = 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 = self._prepare_attn_mask(
             attention_mask,