You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Hi, thanks for your awesome demo of speculative sample. Some of your code maybe outdated in new version of transformer.
In the KVCacheModel class, Bloom model' k cache shape is [bs * head, head_dim, seq] v cache shape is [bs*head, seq, head_dim]
but in transformers 4.44.2, both kv cache shape are (bs, head, seq, head_dim), so this function doesn't work any more.
@torch.no_grad()defrollback(self, end_pos : int):
past_key_values_trimmed= []
assertself._past_key_valuesforkvinself._past_key_values:
k, v=kv# NOTE() the indexing is specific for bloom. This won't work for other models# For example llama k, v should be (batch, num_head, seq_len, hidden_dim)# Bloom is special oneifisinstance(self._model, BloomForCausalLM):
# k (batch * head, hidden_dim, seq); v (batch * head, seq, hidden_dim)k=k[:, :, :end_pos]
v=v[:, :end_pos, :]
kv_trimmed= (k, v)
past_key_values_trimmed.append(kv_trimmed)
else:
# k, v (batch, head, seq, hidden_dim)k=k[:, :, :end_pos, :]
v=v[:, :, :end_pos, :]
kv_trimmed= (k, v)
past_key_values_trimmed.append(kv_trimmed)
self._past_key_values=past_key_values_trimmedself._prob_history=self._prob_history[:, :end_pos, :]
Here is my debug information:
The text was updated successfully, but these errors were encountered:
@shadow150519 Very thankful for pointing out the version incompatibility. Could you please give a Merge Request to this repo? I wish to recognize your contribution here.
Hi, thanks for your awesome demo of speculative sample. Some of your code maybe outdated in new version of transformer.
In the
KVCacheModel
class, Bloom model' k cache shape is [bs * head, head_dim, seq] v cache shape is [bs*head, seq, head_dim]but in transformers 4.44.2, both kv cache shape are (bs, head, seq, head_dim), so this function doesn't work any more.
Here is my debug information:
The text was updated successfully, but these errors were encountered: