Upload modeling_flash_llama.py

modeling_flash_llama.py

@@ -290,9 +290,10 @@ class LlamaAttention(nn.Module):
             scaling_type = self.config.rope_scaling["type"]
             scaling_factor = self.config.rope_scaling["factor"]
             assert scaling_type == 'linear'
-        
+        rotary_base = self.config.__dict__.get("rope_theta", 10000.0)
+
         self.rotary_emb = FlashRotaryEmbedding(
-            self.head_dim, base=10000, interleaved=False, scaling_factor=scaling_factor,
+            self.head_dim, base=rotary_base, interleaved=False, scaling_factor=scaling_factor,
         )
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
@@ -362,12 +363,17 @@ class LlamaAttention(nn.Module):
         past_key_value = (past_kv, past_len+q.size(1)) if use_cache else None
 
         # no padding tokens, more efficient
+        # the basic problem here is that for qlora, stuff is stored in float32, but attention needs float16 or bfloat16.
+        # if you cast just based on torch.cuda.is_bf_16_supported(), it works for training, but it breaks in evals
+        # that load the model in fp16 on a gpu where bf16 is supported. so, we cast based on the GPU support, but then
+        # cast back to whatever q originally was. hopefully that works!
+        orig_dtype = q.dtype
         attn_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
         attn_outputs = flash_attn_kvpacked_func(
             q.type(attn_dtype), kv.type(attn_dtype), dropout_p=0.0, softmax_scale=1.0/self.norm_factor, causal=(not has_layer_past), return_attn_probs=output_attentions)
 
         attn_output = attn_outputs[0] if output_attentions else attn_outputs
-        attn_output = attn_output.reshape(bsz, q_len, h_size)
+        attn_output = attn_output.reshape(bsz, q_len, h_size).type(orig_dtype)
         attn_weights = attn_outputs[2] if output_attentions else None
 
         if self.config.pretraining_tp > 1: