Fix batched inference (#298)
* 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 * Fix inference * swiglu * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update mistral.py * Update llama.py * fast inference * model * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update utils.py * Update llama.py * Update utils.py * inference * Update llama.py * Update llama.py * Update llama.py * overhead * Update llama.py * Update llama.py * compile * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py
This commit is contained in:
Daniel Han
GitHub
parent
8921867157
commit
537577720c
4 changed files with 153 additions and 129 deletions
|
|
@ -118,7 +118,7 @@ def fast_gemv(X, W, quant_state, out = None):
|
|||
if quant_state is None: return torch.matmul(X, W, out = out)
|
||||
# For fast X @ W where seq_len == 1
|
||||
# From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
|
||||
bsz, q_len, hd = X.shape
|
||||
_, q_len, hd = X.shape
|
||||
# assert(q_len == 1)
|
||||
|
||||
if type(quant_state) is not list:
|
||||
|
|
@ -142,10 +142,10 @@ def fast_gemv(X, W, quant_state, out = None):
|
|||
bout = shape[0]
|
||||
|
||||
if out is None:
|
||||
out = torch.empty((bsz, 1, bout,), dtype = dtype, device = "cuda")
|
||||
else:
|
||||
assert(out.shape == (bsz, 1, bout,))
|
||||
pass
|
||||
out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda")
|
||||
# else:
|
||||
# assert(out.shape == (1, 1, bout,))
|
||||
# pass
|
||||
|
||||
n = 1
|
||||
m = shape[0]
|
||||
|
|
@ -171,15 +171,9 @@ def fast_gemv(X, W, quant_state, out = None):
|
|||
fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \
|
||||
cgemm_4bit_inference_naive_bf16
|
||||
|
||||
ptr_W = get_ptr(W)
|
||||
ptr_absmax = get_ptr(absmax)
|
||||
ptr_stats = get_ptr(stats)
|
||||
blocksize = ctypes.c_int32(blocksize)
|
||||
|
||||
for row in range(bsz):
|
||||
fx(m, n, k, get_ptr(X[row]), ptr_W, ptr_absmax, ptr_stats, get_ptr(out[row]),
|
||||
lda, ldb, ldc, blocksize)
|
||||
pass
|
||||
blocksize = ctypes.c_int32(blocksize)
|
||||
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
|
||||
lda, ldb, ldc, blocksize)
|
||||
|
||||
return out
|
||||
pass
|
||||
|
|
@ -189,12 +183,11 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
|
|||
|
||||
W, W_quant, lora_A, lora_B, lora_S = get_lora_parameters(proj)
|
||||
|
||||
bsz, _, in_dim = X.shape
|
||||
bsz, q_len, in_dim = X.shape
|
||||
|
||||
if W_quant is None:
|
||||
out = torch.matmul(X, W.t(), out = out)
|
||||
elif bsz <= 2:
|
||||
# Only batches of 2 are faster with Gemv
|
||||
elif bsz == 1 and q_len == 1:
|
||||
out = fast_gemv(X, W, W_quant, out = out)
|
||||
else:
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
|
|
|
|||
|
|
@ -44,11 +44,11 @@ def fast_geglu_inference(self, X):
|
|||
# gate = self.gate_proj(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")
|
||||
# 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)#, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X)#, out = temp[1])
|
||||
gate = torch_nn_functional_gelu(gate, approximate = "tanh")
|
||||
gate *= up
|
||||
|
||||
|
|
@ -58,18 +58,6 @@ def fast_geglu_inference(self, X):
|
|||
pass
|
||||
|
||||
|
||||
def fast_rms_layernorm_inference_gemma(self, X, out_weight):
|
||||
XX = X.to(torch.float32)
|
||||
variance = XX.square().mean(-1, keepdim = True)
|
||||
variance += self.variance_epsilon
|
||||
XX *= variance.rsqrt_()
|
||||
out_weight[:] = self.weight
|
||||
out_weight += 1.0
|
||||
XX *= out_weight
|
||||
return XX.to(X.dtype)
|
||||
pass
|
||||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590
|
||||
def GemmaDecoderLayer_fast_forward(
|
||||
self,
|
||||
|
|
@ -83,19 +71,21 @@ def GemmaDecoderLayer_fast_forward(
|
|||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
):
|
||||
if past_key_value is not None:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
if use_cache: #past_key_value is not None:
|
||||
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 += residual
|
||||
|
||||
|
|
@ -107,7 +97,6 @@ def GemmaDecoderLayer_fast_forward(
|
|||
else:
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, gemma = True)
|
||||
# hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
causal_mask=causal_mask,
|
||||
|
|
@ -123,19 +112,13 @@ def GemmaDecoderLayer_fast_forward(
|
|||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
|
||||
# hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
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
|
||||
|
||||
|
|
@ -148,31 +131,37 @@ def GemmaModel_fast_forward_inference(
|
|||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
):
|
||||
# Fix out of bounds tokenization
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
out_weight = torch.empty_like(self.layers[0].input_layernorm.weight, dtype = torch.float32, device = "cuda")
|
||||
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
hidden_states = self.model.embed_tokens(input_ids)
|
||||
hidden_states = hidden_states.to(self.config.torch_dtype)
|
||||
# 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32
|
||||
# 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32
|
||||
hidden_states *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = hidden_states.dtype)
|
||||
|
||||
bsz, q_len, hd = hidden_states.shape
|
||||
seq_len = past_key_values[0][0].shape[-2]
|
||||
if bsz != 1:
|
||||
attention_mask = _prepare_4d_causal_attention_mask(attention_mask, (bsz, q_len), hidden_states, seq_len,)
|
||||
pass
|
||||
|
||||
next_decoder_cache = []
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
# Self Attention
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.input_layernorm, hidden_states, out_weight)
|
||||
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
decoder_layer.self_attn,
|
||||
hidden_states,
|
||||
past_key_values[idx],
|
||||
None,
|
||||
hidden_states = hidden_states,
|
||||
past_key_value = past_key_values[idx],
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
)
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.post_attention_layernorm, hidden_states, out_weight)
|
||||
hidden_states = fast_geglu_inference(decoder_layer.mlp, hidden_states)
|
||||
|
|
@ -180,13 +169,13 @@ def GemmaModel_fast_forward_inference(
|
|||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
pass
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(self.norm, hidden_states, out_weight)
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(self.model.norm, hidden_states, out_weight)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state = hidden_states,
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
)
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -74,7 +74,8 @@ pass
|
|||
|
||||
|
||||
from math import sqrt as math_sqrt
|
||||
KV_CACHE_INCREMENT = 128 # KV Cache update size
|
||||
KV_CACHE_INCREMENT = 256 # KV Cache update size
|
||||
torch_nn_functional_softmax = torch.nn.functional.softmax
|
||||
|
||||
def LlamaAttention_fast_forward_inference(
|
||||
self,
|
||||
|
|
@ -82,6 +83,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
|
||||
|
|
@ -138,6 +140,7 @@ def LlamaAttention_fast_forward_inference(
|
|||
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda")
|
||||
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda")
|
||||
self.scalar = 1.0 / math_sqrt(self.head_dim)
|
||||
self.half_head_dim = head_dim // 2
|
||||
elif kv_seq_len >= self.paged_attention.shape[0]:
|
||||
self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim))
|
||||
self.paged_attention_K = self.paged_attention[:,0]
|
||||
|
|
@ -154,17 +157,23 @@ def LlamaAttention_fast_forward_inference(
|
|||
|
||||
# cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
|
||||
# Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
|
||||
cos = self.rotary_emb.cos_cached[seq_len]
|
||||
sin = self.rotary_emb.sin_cached[seq_len]
|
||||
h = head_dim // 2
|
||||
cos = self.rotary_emb.cos_cached[position_ids].unsqueeze(1)
|
||||
sin = self.rotary_emb.sin_cached[position_ids].unsqueeze(1)
|
||||
h = self.half_head_dim
|
||||
|
||||
RH_Q = self.RH_Q
|
||||
RH_Q[:,:,:,:h] = Qn[:,:,:,h:]; RH_Q[:,:,:,h:] = Qn[:,:,:,:h]; torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]);
|
||||
Qn *= cos; Qn.addcmul_(RH_Q, sin);
|
||||
RH_Q[:,:,:,:h] = Qn[:,:,:,h:]
|
||||
RH_Q[:,:,:,h:] = Qn[:,:,:,:h]
|
||||
torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h])
|
||||
Qn *= cos
|
||||
Qn.addcmul_(RH_Q, sin)
|
||||
|
||||
RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda")
|
||||
RH_K[:,:,:,:h] = Kn[:,:,:,h:]; RH_K[:,:,:,h:] = Kn[:,:,:,:h]; torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]);
|
||||
Kn *= cos; Kn.addcmul_(RH_K, sin);
|
||||
RH_K[:,:,:,:h] = Kn[:,:,:,h:]
|
||||
RH_K[:,:,:,h:] = Kn[:,:,:,:h]
|
||||
torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h])
|
||||
Kn *= cos
|
||||
Kn.addcmul_(RH_K, sin)
|
||||
|
||||
# New KV cache
|
||||
# Kn = torch.cat([K1, Kn], dim = 2)
|
||||
|
|
@ -198,10 +207,15 @@ def LlamaAttention_fast_forward_inference(
|
|||
# pass
|
||||
|
||||
# Attention
|
||||
A = torch.matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len])
|
||||
A *= self.scalar
|
||||
A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
|
||||
A = torch.matmul(A, Vnn, out = Qn)
|
||||
if bsz == 1:
|
||||
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)
|
||||
else:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False)
|
||||
pass
|
||||
A = A.transpose(1, 2)
|
||||
A = A.reshape(bsz, 1, attention_size)
|
||||
A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1][:,:,:self.hidden_size])
|
||||
|
|
@ -214,11 +228,11 @@ def fast_swiglu_inference(self, X):
|
|||
# gate = self.gate_proj(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")
|
||||
# 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)#, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X)#, out = temp[1])
|
||||
gate = torch_nn_functional_silu(gate, inplace = True)
|
||||
gate *= up
|
||||
|
||||
|
|
@ -240,6 +254,24 @@ def fast_rms_layernorm_inference(self, X):
|
|||
pass
|
||||
|
||||
|
||||
def fast_rms_layernorm_inference_gemma(self, X, out_weight = None):
|
||||
XX = X.to(torch.float32)
|
||||
variance = XX.square().mean(-1, keepdim = True)
|
||||
variance += self.variance_epsilon
|
||||
XX *= variance.rsqrt_()
|
||||
|
||||
if out_weight is None:
|
||||
out_weight = self.weight + 1.0
|
||||
else:
|
||||
out_weight[:] = self.weight
|
||||
out_weight += 1.0
|
||||
pass
|
||||
|
||||
XX *= out_weight
|
||||
return XX.to(X.dtype)
|
||||
pass
|
||||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L320
|
||||
def LlamaAttention_fast_forward(
|
||||
self,
|
||||
|
|
@ -375,18 +407,18 @@ 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:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
|
||||
# Self Attention
|
||||
if use_cache:
|
||||
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 += residual
|
||||
|
||||
|
|
@ -418,13 +450,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 +629,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,23 +662,24 @@ 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,)
|
||||
# Final layernorm
|
||||
if use_cache:
|
||||
hidden_states = (fast_rms_layernorm_inference_gemma if IS_GEMMA else fast_rms_layernorm_inference)\
|
||||
(self.norm, hidden_states)
|
||||
else:
|
||||
hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA)
|
||||
pass
|
||||
|
||||
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)
|
||||
return BaseModelOutputWithPast(
|
||||
|
|
@ -665,32 +692,44 @@ pass
|
|||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825
|
||||
# @torch.inference_mode
|
||||
def LlamaModel_fast_forward_inference(
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids,
|
||||
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 = self.model.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]
|
||||
if bsz != 1:
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(bsz, q_len),
|
||||
hidden_states,
|
||||
seq_len,
|
||||
sliding_window = getattr(self.config, "sliding_window", None),
|
||||
)
|
||||
else:
|
||||
attention_mask = None
|
||||
pass
|
||||
|
||||
next_decoder_cache = []
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
# Self Attention
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(decoder_layer.input_layernorm, hidden_states)
|
||||
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
decoder_layer.self_attn,
|
||||
hidden_states,
|
||||
past_key_values[idx],
|
||||
None,
|
||||
hidden_states = hidden_states,
|
||||
past_key_value = past_key_values[idx],
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
)
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(decoder_layer.post_attention_layernorm, hidden_states)
|
||||
hidden_states = fast_swiglu_inference(decoder_layer.mlp, hidden_states)
|
||||
|
|
@ -698,13 +737,13 @@ def LlamaModel_fast_forward_inference(
|
|||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
pass
|
||||
hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states)
|
||||
hidden_states = fast_rms_layernorm_inference(self.model.norm, hidden_states)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state = hidden_states,
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
)
|
||||
pass
|
||||
|
||||
|
|
@ -726,11 +765,13 @@ 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 past_key_values is not None:
|
||||
outputs = fast_forward_inference(
|
||||
self.model,
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids = position_ids,
|
||||
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 \
|
||||
hasattr(self.model.layers[0].self_attn, "paged_attention"):
|
||||
if past_key_values is not None:
|
||||
outputs = LlamaModel_fast_forward_inference(
|
||||
self.model,
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
)
|
||||
else:
|
||||
outputs = self.model(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue