Upload modeling_flash_llama.py

modeling_flash_llama.py

@@ -363,7 +363,7 @@ class LlamaAttention(nn.Module):
         past_key_value = (past_kv, past_len+q.size(1)) if use_cache else None
 
         # no padding tokens, more efficient
-        attn_dtype = self.o_proj.weight.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)