Update transformer.py

bigscience-workshop · Sep 16, 2021 · 85c73f2 · 85c73f2
1 parent 7ad5b48
commit 85c73f2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -687,7 +687,7 @@ def build_layer(layer_number):
         if args.position_embedding_type == PositionEmbeddingType.alibi:
             self.alibi = self._build_alibi_tensor(args.seq_length, args.num_attention_heads, args.micro_batch_size).to(torch.cuda.current_device())
             if args.params_dtype == torch.float16:
-                self.alibi = self.alibi.half()
+                self.alibi = self.alibi.to(torch.bfloat16)
         else:
             self.alibi = None
 
@@ -794,4 +794,3 @@ def forward(self, hidden_states, attention_mask, layer_past=None,
             output = [output, presents]
 
         return output
-