Upload modeling_flash_llama.py

modeling_flash_llama.py

@@ -363,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: