Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Fix] Fix the Error of q, k, and v states must have the same dtype when using flash attention forward. #15

Merged
merged 2 commits into from
May 30, 2024
Merged

[Fix] Fix the Error of q, k, and v states must have the same dtype when using flash attention forward. #15

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aba6b66
🚑 Fix: cast back the dtype of q, k, and v to the same target dtype.
jacklanda May 28, 2024
21cd706
✨ Add dtype checking for q, k, and v states before auto casting.
jacklanda May 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions mergoo/models/modeling_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -541,6 +541,26 @@ def _flash_attention_forward(
cu_seqlens_q, cu_seqlens_k = cu_seq_lens
max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens

input_dtype = query_states.dtype
if input_dtype == torch.float32:
jacklanda marked this conversation as resolved.
Show resolved Hide resolved
if torch.is_autocast_enabled():
target_dtype = torch.get_autocast_gpu_dtype()
# Handle the case where the model is quantized
elif hasattr(self.config, "_pre_quantization_dtype"):
target_dtype = self.config._pre_quantization_dtype
else:
target_dtype = self.q_proj.weight.dtype

logger.warning_once(
f"The input hidden states seems to be silently casted in float32, this might be related to"
f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
f" {target_dtype}."
)

query_states = query_states.to(target_dtype)
key_states = key_states.to(target_dtype)
value_states = value_states.to(target_dtype)

attn_output_unpad = flash_attn_varlen_func(
query_states,
key_states,
Expand Down
Loading