Temp fix batch inference (#294)
* Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * llama * Update llama.py * gemma * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update save.py * RoPE * Update llama.py * Update llama.py * Update llama.py * Update gemma.py * correct_dtype * Update gemma.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Chat Templates * Update README.md * Update README.md * Update llama.py * DoRA * Update _utils.py * Update chat_templates.py * Update llama.py * Hotfix - fix DoRA, Gemma prompt template (#202) (#203) * Update save.py * saving * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update __init__.py * Update save.py * Update save.py * Update save.py * save * trainer * spaces * original * Gemma * Update pyproject.toml * Update mapper.py * Update fast_lora.py * FastGemmaModel * model_type * Update llama.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update fast_lora.py * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * gemma * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update fast_lora.py * Update fast_lora.py * Fast CE Loss * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * CE * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update geglu.py * Update cross_entropy_loss.py * revert * Update llama.py * Update llama.py * norm * Update gemma.py * Update gemma.py * position_ids * Update gemma.py * Update gemma.py * pos * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * revert * revert * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * rope * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * llama * Update llama.py * gemma * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update save.py * RoPE * Update llama.py * Update llama.py * Update llama.py * Update gemma.py * correct_dtype * Update gemma.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Chat Templates * Update README.md * Update README.md * Update llama.py * DoRA * Update _utils.py * Update chat_templates.py * Update pyproject.toml * Small fixes * Update pyproject.toml * Approx gelu * Update geglu.py * Approx gelu * Update llama.py * Update __init__.py * Update __init__.py * Update _utils.py * Update geglu.py * Update gemma.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Fix Gemma merging * Update rms_layernorm.py * Update gemma.py * Update pyproject.toml * Layernorms * Gemma precision * Update gemma.py * sqrt * Update gemma.py * Update save.py * RoPE and Gemma precision * Update rms_layernorm.py * Fix warning * Update chat_templates.py * Update chat_templates.py * Update save.py * Update save.py * Update save.py * Update chat_templates.py * Update llama.py * model_name * Update loader.py * Tokenizer overwritten * Update llama.py * Update llama.py * Update llama.py * Update save.py * Accuracy * Revert * Update save.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update chat_templates.py * Update save.py * Update save.py * Update llama.py * Update llama.py * Account for DoRA * Update llama.py * Update save.py * GGUF incorrect * Update save.py * Update pyproject.toml * kaggle new * Update pyproject.toml * Update pyproject.toml * upcasting * Fix Colab * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update rope_embedding.py * Update rope_embedding.py * Fix bugs * Update fast_lora.py * Update fast_lora.py * Update README.md * Update README.md * GGUF * Update save.py * Update save.py * Update save.py * Update save.py * Update README.md * Update README.md * Bugs * Update fast_lora.py * Update pyproject.toml * Update fast_lora.py * Update __init__.py * Update fast_lora.py * dtype * Update llama.py * Update llama.py * Update llama.py * dtype * Update mistral.py * trust_remote_code * lm_head * Update llama.py * save_pretrained_settings * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * state_dict * Update save.py * whoami * Update llama.py * Update save.py * Update llama.py * Patch tokenizer * Update chat_templates.py * Heal tokenizers * Update chat_templates.py * Update mapper.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update chat_templates.py * tokenizer patching * patch_tokenizer * Update chat_templates.py * Update tokenizer_utils.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update tokenizer_utils.py * Edit * Update mistral.py * Update mistral.py * Stats * Update mistral.py * attention_mask * Update llama.py * Update llama.py * batch * Temp fix batch inference * Update llama.py * Update gemma.py
This commit is contained in:
Daniel Han
GitHub
parent
8e263b8b7d
commit
e209991ba1
3 changed files with 81 additions and 57 deletions
|
|
@ -45,10 +45,9 @@ def fast_geglu_inference(self, X):
|
|||
# up = self.up_proj(X)
|
||||
bsz, _, hd = X.shape
|
||||
mlp_size = self.config.intermediate_size
|
||||
temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda")
|
||||
|
||||
gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X, out = temp[1])
|
||||
gate = fast_linear_forward(self.gate_proj, X)
|
||||
up = fast_linear_forward(self. up_proj, X)
|
||||
gate = torch_nn_functional_gelu(gate, approximate = "tanh")
|
||||
gate *= up
|
||||
|
||||
|
|
@ -83,20 +82,30 @@ def GemmaDecoderLayer_fast_forward(
|
|||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
):
|
||||
if past_key_value is not None:
|
||||
if use_cache: #past_key_value is not None:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
out_weight = torch.empty(self.input_layernorm.weight.shape, dtype = torch.float32, device = "cuda")
|
||||
|
||||
# Self Attention
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(self.input_layernorm, hidden_states, out_weight)
|
||||
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
self.self_attn,
|
||||
hidden_states,
|
||||
past_key_value,
|
||||
position_ids,
|
||||
do_prefill = do_prefill,
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
causal_mask=causal_mask,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
padding_mask=padding_mask,
|
||||
)
|
||||
# hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
# self.self_attn,
|
||||
# hidden_states,
|
||||
# past_key_value,
|
||||
# position_ids,
|
||||
# do_prefill = do_prefill,
|
||||
# )
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
|
|
@ -129,13 +138,8 @@ def GemmaDecoderLayer_fast_forward(
|
|||
pass
|
||||
|
||||
outputs = (hidden_states,)
|
||||
|
||||
if output_attentions:
|
||||
outputs += (self_attn_weights,)
|
||||
|
||||
if use_cache:
|
||||
outputs += (present_key_value,)
|
||||
|
||||
if output_attentions: outputs += (self_attn_weights,)
|
||||
if use_cache: outputs += (present_key_value,)
|
||||
return outputs
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ pass
|
|||
|
||||
|
||||
from math import sqrt as math_sqrt
|
||||
KV_CACHE_INCREMENT = 128 # KV Cache update size
|
||||
KV_CACHE_INCREMENT = 256 # KV Cache update size
|
||||
|
||||
def LlamaAttention_fast_forward_inference(
|
||||
self,
|
||||
|
|
@ -82,6 +82,7 @@ def LlamaAttention_fast_forward_inference(
|
|||
past_key_value: Optional[Tuple[torch.Tensor]],
|
||||
position_ids,
|
||||
do_prefill = False,
|
||||
attention_mask = None,
|
||||
):
|
||||
"""
|
||||
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
|
||||
|
|
@ -167,12 +168,12 @@ def LlamaAttention_fast_forward_inference(
|
|||
Kn *= cos; Kn.addcmul_(RH_K, sin);
|
||||
|
||||
# New KV cache
|
||||
# Kn = torch.cat([K1, Kn], dim = 2)
|
||||
# Vn = torch.cat([V1, Vn], dim = 2)
|
||||
self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3)
|
||||
self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)
|
||||
Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
Kn = torch.cat([K1, Kn], dim = 2)
|
||||
Vn = torch.cat([V1, Vn], dim = 2)
|
||||
# self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3)
|
||||
# self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)
|
||||
# Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
# Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
|
||||
# Handle sliding windows
|
||||
sliding_window = getattr(self.config, "sliding_window", None)
|
||||
|
|
@ -200,6 +201,7 @@ def LlamaAttention_fast_forward_inference(
|
|||
# Attention
|
||||
A = torch.matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len])
|
||||
A *= self.scalar
|
||||
if attention_mask is not None: A += attention_mask # Must add attention_mask for batched
|
||||
A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
|
||||
A = torch.matmul(A, Vnn, out = Qn)
|
||||
A = A.transpose(1, 2)
|
||||
|
|
@ -215,10 +217,9 @@ def fast_swiglu_inference(self, X):
|
|||
# up = self.up_proj(X)
|
||||
bsz, _, hd = X.shape
|
||||
mlp_size = self.config.intermediate_size
|
||||
temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda")
|
||||
|
||||
gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X, out = temp[1])
|
||||
gate = fast_linear_forward(self.gate_proj, X)
|
||||
up = fast_linear_forward(self. up_proj, X)
|
||||
gate = torch_nn_functional_silu(gate, inplace = True)
|
||||
gate *= up
|
||||
|
||||
|
|
@ -375,19 +376,30 @@ def LlamaDecoderLayer_fast_forward(
|
|||
(see `past_key_values`).
|
||||
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
|
||||
"""
|
||||
if past_key_value is not None:
|
||||
if use_cache: #past_key_value is not None:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
|
||||
# Self Attention
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
|
||||
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
self.self_attn,
|
||||
hidden_states,
|
||||
past_key_value,
|
||||
position_ids,
|
||||
do_prefill = do_prefill,
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
causal_mask=causal_mask,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
padding_mask=padding_mask,
|
||||
)
|
||||
# hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
# self.self_attn,
|
||||
# hidden_states,
|
||||
# past_key_value,
|
||||
# position_ids,
|
||||
# do_prefill = do_prefill,
|
||||
# attention_mask = attention_mask,
|
||||
# )
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
|
|
@ -418,13 +430,8 @@ def LlamaDecoderLayer_fast_forward(
|
|||
pass
|
||||
|
||||
outputs = (hidden_states,)
|
||||
|
||||
if output_attentions:
|
||||
outputs += (self_attn_weights,)
|
||||
|
||||
if use_cache:
|
||||
outputs += (present_key_value,)
|
||||
|
||||
if output_attentions: outputs += (self_attn_weights,)
|
||||
if use_cache: outputs += (present_key_value,)
|
||||
return outputs
|
||||
pass
|
||||
|
||||
|
|
@ -602,9 +609,8 @@ def LlamaModel_fast_forward(
|
|||
pass
|
||||
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if output_hidden_states: all_hidden_states += (hidden_states,)
|
||||
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
|
|
@ -636,22 +642,15 @@ def LlamaModel_fast_forward(
|
|||
use_cache=use_cache,
|
||||
padding_mask=padding_mask,
|
||||
)
|
||||
pass
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if use_cache:
|
||||
next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
|
||||
|
||||
if output_attentions:
|
||||
all_self_attns += (layer_outputs[1],)
|
||||
if use_cache: next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
|
||||
if output_attentions: all_self_attns += (layer_outputs[1],)
|
||||
pass
|
||||
|
||||
hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if output_hidden_states: all_hidden_states += (hidden_states,)
|
||||
next_cache = next_decoder_cache if use_cache else None
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
||||
|
|
@ -670,12 +669,29 @@ def LlamaModel_fast_forward_inference(
|
|||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
attention_mask = None,
|
||||
):
|
||||
# Fix out of bounds tokenization
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
hidden_states = hidden_states.to(self.config.torch_dtype)
|
||||
bsz, q_len, hd = hidden_states.shape
|
||||
seq_len = past_key_values[0][0].shape[-2]
|
||||
|
||||
# Must use attention mask for batched processing
|
||||
sliding_window = getattr(self.config, "sliding_window", None)
|
||||
if (sliding_window is not None and seq_len >= sliding_window) or (bsz != 1):
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(bsz, q_len),
|
||||
hidden_states,
|
||||
seq_len,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
else:
|
||||
attention_mask = None
|
||||
pass
|
||||
|
||||
next_decoder_cache = []
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
|
|
@ -686,7 +702,9 @@ def LlamaModel_fast_forward_inference(
|
|||
decoder_layer.self_attn,
|
||||
hidden_states,
|
||||
past_key_values[idx],
|
||||
None,
|
||||
position_ids = None,
|
||||
do_prefill = False,
|
||||
attention_mask = attention_mask,
|
||||
)
|
||||
hidden_states += residual
|
||||
|
||||
|
|
@ -726,11 +744,12 @@ def CausalLM_fast_forward(fast_forward_inference):
|
|||
*args, **kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
|
||||
if past_key_values is not None and hasattr(self.model.layers[0].self_attn, "paged_attention"):
|
||||
if False:#past_key_values is not None and hasattr(self.model.layers[0].self_attn, "paged_attention"):
|
||||
outputs = fast_forward_inference(
|
||||
self.model,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
attention_mask = attention_mask,
|
||||
)
|
||||
else:
|
||||
causal_mask = xformers.attn_bias.LowerTriangularMask()
|
||||
|
|
|
|||
|
|
@ -200,12 +200,13 @@ def MistralForCausalLM_fast_forward(
|
|||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
self.model._has_no_labels = labels is None
|
||||
|
||||
if past_key_values is not None and \
|
||||
if False and past_key_values is not None and \
|
||||
hasattr(self.model.layers[0].self_attn, "paged_attention"):
|
||||
outputs = LlamaModel_fast_forward_inference(
|
||||
self.model,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
attention_mask = attention_mask,
|
||||
)
|
||||
else:
|
||||
outputs = self.model(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue