Update transformer.py

EleutherAI · Nov 27, 2023 · f745997 · f745997
1 parent fc214d0
commit f745997
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -658,12 +658,13 @@ def forward(self, hidden_states, attention_mask, layer_past=None):
                 # full rotary
                 query_rot, key_rot = query_layer, key_layer
 
-            if self.bf16:
-                apply_rotary_fn = apply_rotary_pos_emb_torch
-            elif self.rope_fusion:
+            if self.rope_fusion:
                 apply_rotary_fn = fused_apply_rotary_pos_emb
-            else:
-                apply_rotary_fn = apply_rotary_pos_emb
+            else:            
+                if self.bf16:
+                    apply_rotary_fn = apply_rotary_pos_emb_torch
+                else:
+                    apply_rotary_fn = apply_rotary_pos_emb
 
             seq_len = key_layer.shape[0]
             offset = 0