Update transformer.py

Just checked and this should work for bf16. Or, at least, the reason I originally thought it wouldn't doesn't apply.
EleutherAI · Nov 27, 2023 · fc214d0 · fc214d0
1 parent e35955b
commit fc214d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -660,7 +660,7 @@ def forward(self, hidden_states, attention_mask, layer_past=None):
 
             if self.bf16:
                 apply_rotary_fn = apply_rotary_pos_emb_torch
-            elif self.rope_fusion: # we do not currently support bf16 + fused rotary
+            elif self.rope_fusion:
                 apply_rotary_fn = fused_apply_rotary_pos_emb
             else:
                 apply_rotary_fn = apply_rotary_pos_emb