github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Hotfix] Fix bugs in testing continuous batching (#5270)

Browse Source 
* fix bug

* fix bugs

* fix bugs

* fix bugs and add padding

* add funcs and fix bugs

* fix typos

* fix bugs

* add func
This commit is contained in:
Jianghai
2024-01-18 16:31:14 +08:00
committed by GitHub
parent 5ae9099f92
commit 9e2342bde2
6 changed files with 86 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/inference/modeling/models/llama.py
View File

@@ -173,7 +173,10 @@ def llama_attn_forward(
key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
kv_seq_len = max(sequence_lengths).item()
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
query_states = query_states.transpose(1, 2)