Update gemma.py
This commit is contained in:
Daniel Han-Chen
parent
7a5db6758c
commit
1f450cd6ee
1 changed files with 1 additions and 389 deletions
|
|
@ -57,173 +57,6 @@ def fast_geglu_inference(self, X):
|
|||
pass
|
||||
|
||||
|
||||
class FastGemmaRotaryEmbedding(torch.nn.Module):
|
||||
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.base = base
|
||||
self.register_buffer("inv_freq", None, persistent=False)
|
||||
self.register_buffer("cos_cached", None, persistent=False)
|
||||
self.register_buffer("sin_cached", None, persistent=False)
|
||||
|
||||
self.inv_freq = 1.0 / (
|
||||
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cuda").float() / self.dim)
|
||||
)
|
||||
|
||||
position_ids = torch.arange(self.max_position_embeddings, device="cuda", dtype=torch.int64).unsqueeze(0)
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
self.cos_cached2 = freqs
|
||||
self.sin_cached2 = emb.sin().to(torch.bfloat16)
|
||||
|
||||
def forward(self, x, position_ids, seq_len=None):
|
||||
# x: [bs, num_attention_heads, seq_len, head_size]
|
||||
if self.inv_freq is None:
|
||||
self.inv_freq = 1.0 / (
|
||||
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)
|
||||
)
|
||||
|
||||
# length = position_ids.shape[1]
|
||||
if self.cos_cached is None:
|
||||
position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0)
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
self.cos_cached = emb.cos().to(dtype=x.dtype)
|
||||
self.sin_cached = emb.sin().to(dtype=x.dtype)
|
||||
pass
|
||||
|
||||
position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0)
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos_cached = emb.cos().to(dtype=x.dtype)
|
||||
sin_cached = emb.sin().to(dtype=x.dtype)
|
||||
|
||||
print(freqs)
|
||||
print(self.cos_cached2)
|
||||
|
||||
# return self.cos_cached[:,:length], self.sin_cached[:,:length]
|
||||
return self.cos_cached, self.sin_cached
|
||||
pass
|
||||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L320
|
||||
def GemmaAttention_fast_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value = None, #Optional[Cache] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
):
|
||||
# Clear inference
|
||||
if hasattr(self, "paged_attention"):
|
||||
del self.paged_attention_K
|
||||
del self.paged_attention_V
|
||||
del self.paged_attention
|
||||
del self.temp_QA
|
||||
del self.temp_KV
|
||||
del self.RH_Q
|
||||
del self.attention
|
||||
pass
|
||||
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
n_heads = self.num_heads
|
||||
n_groups = self.num_key_value_groups
|
||||
n_kv_heads = self.num_key_value_heads
|
||||
head_dim = self.head_dim
|
||||
assert(n_kv_heads * n_groups == n_heads)
|
||||
|
||||
Q, K, V = self.apply_qkv(self, hidden_states)
|
||||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
|
||||
if True:# position_ids is None:
|
||||
# cos = self.rotary_emb.cos_cached
|
||||
# sin = self.rotary_emb.sin_cached
|
||||
cos, sin = self.rotary_emb(V, position_ids, seq_len = q_len)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
else:
|
||||
cos, sin = self.rotary_emb(V, position_ids, seq_len = q_len)
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
pass
|
||||
|
||||
past_key_value = getattr(self, "past_key_value", past_key_value)
|
||||
|
||||
if past_key_value is not None:
|
||||
# sin and cos are specific to RoPE models; position_ids needed for the static cache
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
||||
K, V = past_key_value.update(K, V, self.layer_idx, cache_kwargs)
|
||||
|
||||
# Attention module
|
||||
if (not HAS_FLASH_ATTENTION):# and attention_mask is None):
|
||||
# Xformers memory efficient attention
|
||||
# Also has Flash Attention v2 dispatching
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
|
||||
# Group query attention
|
||||
if n_groups != 1:
|
||||
K = K .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
pass
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
A = flash_attn_func(Q, K, V, causal = True)
|
||||
else:
|
||||
causal_mask = attention_mask
|
||||
if attention_mask is not None and cache_position is not None:
|
||||
causal_mask = causal_mask[:, :, cache_position, : K.shape[-2]]
|
||||
|
||||
# Grouped query attention
|
||||
if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(Q, K, V, attn_mask = causal_mask, is_causal = False)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
pass
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
return attn_output, attn_weights, past_key_value
|
||||
pass
|
||||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590
|
||||
def GemmaDecoderLayer_fast_forward(
|
||||
self,
|
||||
|
|
@ -340,227 +173,6 @@ def GemmaModel_fast_forward_inference(
|
|||
pass
|
||||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825
|
||||
def GemmaModel_fast_forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError(
|
||||
"You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
|
||||
)
|
||||
|
||||
if self.gradient_checkpointing and self.training and use_cache:
|
||||
logger.warning_once(
|
||||
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
|
||||
)
|
||||
use_cache = False
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
past_seen_tokens = 0
|
||||
if use_cache: # kept for BC (cache positions)
|
||||
if not isinstance(past_key_values, StaticCache):
|
||||
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
|
||||
past_seen_tokens = past_key_values.get_seq_length()
|
||||
|
||||
if cache_position is None:
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
)
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
causal_mask = self._update_causal_mask(attention_mask, inputs_embeds)
|
||||
|
||||
# embed positions
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
# normalized
|
||||
hidden_states = hidden_states * (self.config.hidden_size**0.5)
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = None
|
||||
|
||||
for decoder_layer in self.layers:
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
decoder_layer.__call__,
|
||||
hidden_states,
|
||||
causal_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
output_attentions,
|
||||
use_cache,
|
||||
cache_position,
|
||||
)
|
||||
else:
|
||||
layer_outputs = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_values,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
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],)
|
||||
|
||||
|
||||
# hidden_states = self.norm(hidden_states)
|
||||
hidden_states = fast_rms_layernorm(self.norm, hidden_states)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
next_cache = None
|
||||
if use_cache:
|
||||
next_cache = (
|
||||
next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache
|
||||
)
|
||||
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(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=next_cache,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attns,
|
||||
)
|
||||
pass
|
||||
|
||||
|
||||
def GemmaForCausalLM_fast_forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
*args, **kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
|
||||
if causal_mask is None and past_key_values is None:
|
||||
causal_mask = xformers.attn_bias.LowerTriangularMask()
|
||||
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# 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"):
|
||||
outputs = GemmaModel_fast_forward_inference(
|
||||
self.model,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
)
|
||||
else:
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
causal_mask=causal_mask,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
pass
|
||||
|
||||
hidden_states = outputs[0]
|
||||
bsz, q_len, hd = hidden_states.shape
|
||||
if bsz == 1 and q_len == 1:
|
||||
logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
|
||||
logits = logits.unsqueeze(0).unsqueeze(0)
|
||||
else:
|
||||
logits = self.lm_head(hidden_states)
|
||||
pass
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
shift_logits = logits
|
||||
if not hasattr(self, "extra_ignored_labels"):
|
||||
# Fixes https://github.com/unslothai/unsloth/issues/10
|
||||
self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda")
|
||||
pass
|
||||
|
||||
shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
|
||||
loss = fast_cross_entropy_loss(
|
||||
logits = shift_logits,
|
||||
labels = shift_labels,
|
||||
)
|
||||
pass
|
||||
# if labels is not None:
|
||||
# # Shift so that tokens < n predict n
|
||||
# shift_logits = logits[..., :-1, :].contiguous()
|
||||
# shift_labels = labels[..., 1:].contiguous()
|
||||
# # Flatten the tokens
|
||||
# loss_fct = torch.nn.CrossEntropyLoss()
|
||||
# shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
# shift_labels = shift_labels.view(-1)
|
||||
# # Enable model parallelism
|
||||
# shift_labels = shift_labels.to(shift_logits.device)
|
||||
# loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
pass
|
||||
|
||||
|
||||
class FastGemmaModel(FastLlamaModel):
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -631,7 +243,7 @@ class FastGemmaModel(FastLlamaModel):
|
|||
# Downcast RoPE embedding to correct data type
|
||||
if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \
|
||||
and (module.cos_cached.dtype != correct_dtype):
|
||||
|
||||
|
||||
module.cos_cached = module.cos_cached.to(correct_dtype)
|
||||
module.sin_cached = module.sin_cached.to(correct_dtype)
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue