初始化项目，由ModelHub XC社区提供模型
Model: celestialcreator/Llama-3.2-1B-MTP-k8 Source: Original Platform
This commit is contained in:
ModelHub XC
969
modeling_llama.py
Normal file
969
modeling_llama.py
Normal file
@@ -0,0 +1,969 @@
|
||||
# fmt: off
|
||||
|
||||
# Adaptation recipe lifted from Jonas et al. :>
|
||||
# https://github.com/seal-rg/recurrent-pretraining/blob/main/recpre/raven_modeling_minimal.py
|
||||
|
||||
# coding=utf-8
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Callable, Optional, Union
|
||||
import time
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
from transformers.generation import GenerationMixin
|
||||
from transformers.integrations import use_kernel_forward_from_hub
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_layers import (
|
||||
GenericForQuestionAnswering,
|
||||
GenericForSequenceClassification,
|
||||
GenericForTokenClassification,
|
||||
GradientCheckpointingLayer,
|
||||
)
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
||||
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
||||
from transformers.processing_utils import Unpack
|
||||
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
|
||||
from transformers.utils.deprecation import deprecate_kwarg
|
||||
from transformers.utils.generic import check_model_inputs
|
||||
from .configuration_llama import LlamaConfig
|
||||
|
||||
|
||||
# Glue
|
||||
from transformers.generation.utils import GenerateDecoderOnlyOutput
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
@use_kernel_forward_from_hub("RMSNorm")
|
||||
class LlamaRMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
LlamaRMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return self.weight * hidden_states.to(input_dtype)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
||||
|
||||
|
||||
class LlamaRotaryEmbedding(nn.Module):
|
||||
inv_freq: torch.Tensor # fix linting for `register_buffer`
|
||||
|
||||
def __init__(self, config: LlamaConfig, device=None):
|
||||
super().__init__()
|
||||
# BC: "rope_type" was originally "type"
|
||||
if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
|
||||
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
|
||||
else:
|
||||
self.rope_type = "default"
|
||||
self.max_seq_len_cached = config.max_position_embeddings
|
||||
self.original_max_seq_len = config.max_position_embeddings
|
||||
|
||||
self.config = config
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.original_inv_freq = self.inv_freq
|
||||
|
||||
@torch.no_grad()
|
||||
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
||||
def forward(self, x, position_ids):
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
|
||||
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
||||
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos = emb.cos() * self.attention_scaling
|
||||
sin = emb.sin() * self.attention_scaling
|
||||
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden dims of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
"""Applies Rotary Position Embedding to the query and key tensors.
|
||||
|
||||
Args:
|
||||
q (`torch.Tensor`): The query tensor.
|
||||
k (`torch.Tensor`): The key tensor.
|
||||
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
||||
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
||||
position_ids (`torch.Tensor`, *optional*):
|
||||
Deprecated and unused.
|
||||
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
||||
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
||||
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
||||
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
||||
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
||||
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
||||
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
||||
Returns:
|
||||
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
||||
"""
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
class LlamaMLP(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x):
|
||||
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
return down_proj
|
||||
|
||||
|
||||
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""
|
||||
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
||||
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
||||
"""
|
||||
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
||||
if n_rep == 1:
|
||||
return hidden_states
|
||||
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
||||
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
||||
|
||||
|
||||
def eager_attention_forward(
|
||||
module: nn.Module,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
scaling: float,
|
||||
dropout: float = 0.0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
):
|
||||
key_states = repeat_kv(key, module.num_key_value_groups)
|
||||
value_states = repeat_kv(value, module.num_key_value_groups)
|
||||
|
||||
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
||||
if attention_mask is not None:
|
||||
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
||||
attn_weights = attn_weights + causal_mask
|
||||
|
||||
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
||||
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
||||
attn_output = torch.matmul(attn_weights, value_states)
|
||||
attn_output = attn_output.transpose(1, 2).contiguous()
|
||||
|
||||
return attn_output, attn_weights
|
||||
|
||||
|
||||
class LlamaAttention(nn.Module):
|
||||
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||
|
||||
def __init__(self, config: LlamaConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer_idx = layer_idx
|
||||
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
||||
self.scaling = self.head_dim**-0.5
|
||||
self.attention_dropout = config.attention_dropout
|
||||
self.is_causal = True
|
||||
|
||||
self.q_proj = nn.Linear(
|
||||
config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.k_proj = nn.Linear(
|
||||
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.v_proj = nn.Linear(
|
||||
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
|
||||
)
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
past_key_values: Optional[Cache] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
input_shape = hidden_states.shape[:-1]
|
||||
hidden_shape = (*input_shape, -1, self.head_dim)
|
||||
|
||||
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||
|
||||
if past_key_values is not None:
|
||||
# sin and cos are specific to RoPE models; cache_position needed for the static cache
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
||||
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||
|
||||
attention_interface: Callable = eager_attention_forward
|
||||
if self.config._attn_implementation != "eager":
|
||||
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
||||
|
||||
attn_output, attn_weights = attention_interface(
|
||||
self,
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attention_mask,
|
||||
dropout=0.0 if not self.training else self.attention_dropout,
|
||||
scaling=self.scaling,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
||||
attn_output = self.o_proj(attn_output)
|
||||
return attn_output, attn_weights
|
||||
|
||||
|
||||
class LlamaDecoderLayer(GradientCheckpointingLayer):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
|
||||
|
||||
self.mlp = LlamaMLP(config)
|
||||
self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
# Self Attention
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaPreTrainedModel(PreTrainedModel):
|
||||
config: LlamaConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["LlamaDecoderLayer"]
|
||||
_skip_keys_device_placement = ["past_key_values"]
|
||||
_supports_flash_attn = True
|
||||
_supports_sdpa = True
|
||||
_supports_flex_attn = True
|
||||
|
||||
_can_compile_fullgraph = True
|
||||
_supports_attention_backend = True
|
||||
_can_record_outputs = {
|
||||
"hidden_states": LlamaDecoderLayer,
|
||||
"attentions": LlamaAttention,
|
||||
}
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaModel(LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.rotary_emb = LlamaRotaryEmbedding(config=config)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@check_model_inputs
|
||||
@auto_docstring
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> BaseModelOutputWithPast:
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
|
||||
|
||||
if use_cache and past_key_values is None:
|
||||
past_key_values = DynamicCache(config=self.config)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position: torch.Tensor = 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 = create_causal_mask(
|
||||
config=self.config,
|
||||
input_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
cache_position=cache_position,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||
|
||||
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=past_key_values,
|
||||
)
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
print(f"Loading local LlamaForCausalLM!")
|
||||
|
||||
@can_return_tuple
|
||||
@auto_docstring
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> CausalLMOutputWithPast:
|
||||
r"""
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoTokenizer, LlamaForCausalLM
|
||||
|
||||
>>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
|
||||
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
||||
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
>>> # Generate
|
||||
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
||||
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
||||
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
||||
```"""
|
||||
outputs: BaseModelOutputWithPast = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids: Optional[torch.LongTensor] = None, *args, **kwargs):
|
||||
"""
|
||||
Custom generate entry point.
|
||||
|
||||
If `do_mtp=True` is passed, strictly enforces MTP-only arguments and routes to `_mtp_generate`.
|
||||
Otherwise, routes to standard Hugging Face `generate`.
|
||||
"""
|
||||
|
||||
# 1. Standard Path
|
||||
if not kwargs.get("do_mtp", False):
|
||||
print("Executing standard generate() codepath.")
|
||||
return super().generate(input_ids, *args, **kwargs)
|
||||
|
||||
# 2. MTP Path
|
||||
print("Executing custom MTP generation codepath!")
|
||||
|
||||
# Handle input_ids: HF can pass it as positional (first arg) or keyword
|
||||
# We consolidate it into 'prompt' for the MTP signature
|
||||
prompt = input_ids if input_ids is not None else kwargs.pop("input_ids", None)
|
||||
if prompt is None:
|
||||
# If standard generate was called without input_ids, it might be in *args or handled deeper,
|
||||
# but for MTP we require it explicitly.
|
||||
raise ValueError("MTP generation requires 'input_ids' to be passed.")
|
||||
|
||||
# --- Argument Extraction & Strict Validation ---
|
||||
|
||||
# Keys strictly allowed for MTP (will be passed to _mtp_generate)
|
||||
mtp_allowed_keys = {
|
||||
"do_mtp", "k_toks", "mask_id", "min_mask_id", "max_mask_id", "strategy", "return_mtp_result_dict", "include_prompt", "streamer"
|
||||
}
|
||||
|
||||
# Keys from HF that we know how to map to MTP equivalents
|
||||
# max_length -> max_returned_tokens
|
||||
max_length = kwargs.pop("max_length", None)
|
||||
max_returned_tokens = kwargs.pop("max_returned_tokens", None)
|
||||
if max_returned_tokens is None and max_length is not None:
|
||||
print(f"Renaming max_length={max_length} to max_returned_tokens for MTP generation.")
|
||||
max_returned_tokens = max_length
|
||||
|
||||
# eos_token_id -> eos_id
|
||||
eos_token_id = kwargs.pop("eos_token_id", None)
|
||||
eos_id = kwargs.pop("eos_id", None)
|
||||
if eos_id is None:
|
||||
eos_id = eos_token_id
|
||||
|
||||
# Standard HF args that we SILENTLY IGNORE because they are passed automatically
|
||||
# by the GenerationMixin but are not relevant or supported in this MTP implementation.
|
||||
ignored_hf_keys = {
|
||||
"attention_mask", "use_cache", "do_sample", "stopping_criteria",
|
||||
"pad_token_id", "logits_processor", "max_new_tokens", "generation_config",
|
||||
}
|
||||
|
||||
# Check for explicit incompatibility
|
||||
if kwargs.get("do_sample", False):
|
||||
raise ValueError("MTP generation does not support sampling (do_sample=True).")
|
||||
|
||||
# Extract valid MTP args
|
||||
mtp_kwargs = {}
|
||||
for k in list(kwargs.keys()):
|
||||
if k in mtp_allowed_keys:
|
||||
mtp_kwargs[k] = kwargs.pop(k)
|
||||
|
||||
# Remove ignored HF keys
|
||||
for k in list(kwargs.keys()):
|
||||
if k in ignored_hf_keys:
|
||||
kwargs.pop(k)
|
||||
|
||||
# FAIL LOUDLY if anything is left in kwargs
|
||||
if kwargs:
|
||||
raise ValueError(
|
||||
f"Unsupported argument(s) passed to MTP generate: {list(kwargs.keys())}.\n"
|
||||
f"When do_mtp=True, only these args are supported: {list(mtp_allowed_keys) + ['max_returned_tokens', 'eos_id']}."
|
||||
)
|
||||
|
||||
# Pre-flight checks
|
||||
if len(prompt.shape) > 1 and prompt.shape[0] > 1:
|
||||
raise NotImplementedError("MTP generation currently only supports single-example generation (no batching).")
|
||||
|
||||
# Execute Unified Implementation
|
||||
# Note: We remove 'do_mtp' from kwargs before passing, as the impl doesn't need it
|
||||
mtp_kwargs.pop("do_mtp", None)
|
||||
|
||||
return self._mtp_generate(
|
||||
prompt=prompt,
|
||||
max_returned_tokens=max_returned_tokens,
|
||||
eos_id=eos_id,
|
||||
**mtp_kwargs
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _mtp_generate(
|
||||
self,
|
||||
prompt: torch.Tensor,
|
||||
max_returned_tokens: int = None,
|
||||
k_toks: int = 1,
|
||||
mask_id: int = None,
|
||||
min_mask_id: int = None,
|
||||
max_mask_id: int = None,
|
||||
eos_id: Optional[Union[int, list]] = None,
|
||||
include_prompt: bool = True,
|
||||
streamer = None,
|
||||
strategy: Optional[list] = None,
|
||||
return_mtp_result_dict: bool = False,
|
||||
):
|
||||
"""
|
||||
Implementation of MTP generation logic.
|
||||
"""
|
||||
# --- Setup Stop Tokens ---
|
||||
if isinstance(eos_id, int):
|
||||
stop_tokens = ([eos_id,],)
|
||||
elif isinstance(eos_id, list):
|
||||
assert all(isinstance(eid, int) for eid in eos_id), "If eos_id is a list, all elements must be ints."
|
||||
stop_tokens = tuple([list([eid,]) for eid in eos_id])
|
||||
elif eos_id is None:
|
||||
stop_tokens = ()
|
||||
else:
|
||||
raise ValueError(f"eos_id must be None, int, or list of lists, got {type(eos_id)}")
|
||||
|
||||
# --- Validation ---
|
||||
if k_toks > 1: assert (mask_id is not None) or (min_mask_id is not None), "mask_id must be provided when k_toks > 1"
|
||||
|
||||
input_ids = prompt.clone()
|
||||
prompt_size = prompt.size(1)
|
||||
device = prompt.device
|
||||
|
||||
# Get generation config (defaulting to model's if not present)
|
||||
generation_config = self.generation_config
|
||||
|
||||
# --- Streaming Prompt ---
|
||||
if include_prompt:
|
||||
if streamer is not None:
|
||||
print(f"\n<BEGIN Streaming Prompt>", flush=True)
|
||||
streamer.put(input_ids)
|
||||
print(f"\n<END Streaming Prompt>", flush=True)
|
||||
|
||||
if streamer is not None:
|
||||
print(f"\n<BEGIN Streaming Generation>", flush=True)
|
||||
|
||||
stop_progress = [0] * len(stop_tokens)
|
||||
|
||||
# --- Generation Loop ---
|
||||
t0_prefill = time.perf_counter()
|
||||
t1_prefill = None
|
||||
t0_gen = None
|
||||
t1_gen = None
|
||||
toks_pre_prefill = input_ids.shape[1]
|
||||
toks_post_prefill = None
|
||||
current_idx = input_ids.shape[1]
|
||||
num_fwd_evals = 0
|
||||
effective_k_values = []
|
||||
|
||||
# Prepare kwargs for the inner loop
|
||||
model_kwargs = {}
|
||||
|
||||
while current_idx + k_toks <= max_returned_tokens:
|
||||
|
||||
# 0 is prefill, 1 is first step which can include compile time, then 2 is steady state
|
||||
if (t0_gen is None and num_fwd_evals == 2):
|
||||
t1_prefill = time.perf_counter()
|
||||
t0_gen = time.perf_counter()
|
||||
toks_post_prefill = input_ids.shape[1]
|
||||
|
||||
# Generate the token
|
||||
if k_toks > 1:
|
||||
input_ids = self._extend_w_mask(input_ids=input_ids, k_toks=k_toks, mask_id=mask_id, min_mask_id=min_mask_id, max_mask_id=max_mask_id)
|
||||
|
||||
# first step prep
|
||||
if num_fwd_evals == 0:
|
||||
model_kwargs, generation_config = self._prep_generate_args(
|
||||
self,
|
||||
input_ids,
|
||||
generation_config,
|
||||
)
|
||||
assert "token_type_ids" not in model_kwargs
|
||||
assert "attention_mask" not in model_kwargs
|
||||
assert "decoder_attention_mask" not in model_kwargs
|
||||
|
||||
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
|
||||
|
||||
next_tokens, model_outputs = self._mtp_next_tokens(
|
||||
self,
|
||||
model_inputs,
|
||||
k_toks=k_toks,
|
||||
strategy=strategy,
|
||||
)
|
||||
effective_k_values.append(next_tokens.shape[1])
|
||||
|
||||
if streamer is not None: streamer.put(next_tokens)
|
||||
|
||||
# remove the masks if any
|
||||
if k_toks > 1:
|
||||
input_ids = input_ids[:, :-(k_toks - 1)]
|
||||
input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||||
|
||||
# Update cache / model kwargs
|
||||
model_kwargs["past_key_values"] = model_outputs.past_key_values
|
||||
|
||||
if strategy is None:
|
||||
if num_fwd_evals == 0:
|
||||
# this is the end of the prefill step
|
||||
if k_toks > 1:
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
prompt_size,
|
||||
prompt_size + k_toks + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"] = torch.tensor(
|
||||
[prompt_size], device=device, dtype=torch.int64
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"].add_(k_toks)
|
||||
else: # we can assume that all strats produce variable number of tokens
|
||||
num_new_tokens = next_tokens.shape[1]
|
||||
if num_fwd_evals == 0:
|
||||
# this is the end of the prefill step
|
||||
if k_toks > 1:
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
prompt_size,
|
||||
prompt_size + num_new_tokens + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"] = torch.tensor(
|
||||
[prompt_size], device=device, dtype=torch.int64
|
||||
)
|
||||
else:
|
||||
if k_toks > 1:
|
||||
recomputation_positions = model_kwargs["cache_position"][: -(k_toks - 1)]
|
||||
previous_num_new_tokens = recomputation_positions.size(0)
|
||||
new_start_pos = recomputation_positions[0] + previous_num_new_tokens
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
new_start_pos,
|
||||
new_start_pos + num_new_tokens + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"].add_(1)
|
||||
|
||||
current_idx += next_tokens.shape[1]
|
||||
num_fwd_evals += 1
|
||||
|
||||
# Crop cache
|
||||
model_kwargs["past_key_values"].crop(model_kwargs["cache_position"][0])
|
||||
|
||||
# Check for stop sequences
|
||||
hit_stop_seq = False
|
||||
for int_tok in next_tokens.tolist()[0]: # assuming batch size 1
|
||||
for i, seq in enumerate(stop_tokens):
|
||||
if int_tok == seq[stop_progress[i]]:
|
||||
stop_progress[i] += 1
|
||||
if stop_progress[i] == len(seq):
|
||||
hit_stop_seq = True
|
||||
break
|
||||
else:
|
||||
stop_progress[i] = 0
|
||||
if hit_stop_seq:
|
||||
break
|
||||
if hit_stop_seq:
|
||||
break
|
||||
|
||||
# End of generation loop
|
||||
if streamer is not None:
|
||||
streamer.end()
|
||||
print(f"\n<END Streaming Generation>", flush=True)
|
||||
|
||||
t1_gen = time.perf_counter()
|
||||
|
||||
# Calculate stats
|
||||
if t1_prefill is not None and t0_gen is not None and toks_post_prefill is not None:
|
||||
t_prefill = t1_prefill - t0_prefill
|
||||
t_gen = t1_gen - t0_gen
|
||||
tokens_generated = input_ids.shape[1] - toks_post_prefill
|
||||
toks_gend_incl_prefillplus1 = input_ids.shape[1] - toks_pre_prefill
|
||||
else:
|
||||
t_prefill = t1_gen - t0_prefill
|
||||
t_gen = t1_gen - t0_prefill
|
||||
tokens_generated = input_ids.shape[1] - toks_pre_prefill
|
||||
toks_gend_incl_prefillplus1 = input_ids.shape[1] - toks_pre_prefill
|
||||
|
||||
print(
|
||||
f"Using a total of {f'1+1+{num_fwd_evals-2}' if t0_gen is not None else f'{num_fwd_evals}'} forward evals, time for prefill plus first/compilation step {f'(1+1)' if t0_gen is not None else ''} was {t_prefill:.02f} sec, generation time was {t_gen:.02f} sec @ {tokens_generated / t_gen:.02f} tokens/sec {f'steady state' if t0_gen is not None else ''} over {tokens_generated} tokens.",
|
||||
flush=True,
|
||||
)
|
||||
print(f"Strategy used: {strategy}", flush=True)
|
||||
avg_effective_k = sum(effective_k_values) / len(effective_k_values) if effective_k_values else 0
|
||||
print(
|
||||
f"Average effective k_toks over generation: {avg_effective_k:.02f}, full array of effective k_toks: {effective_k_values}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if generation_config.return_dict_in_generate:
|
||||
raise NotImplementedError(f"Only basic return type implemented for MTP generate.")
|
||||
|
||||
if return_mtp_result_dict:
|
||||
token_ids = input_ids if include_prompt else input_ids[:,prompt_size:]
|
||||
|
||||
# testing for generation vs. aux data order integrity
|
||||
import hashlib
|
||||
leading_toks = token_ids[0,prompt_size:prompt_size+10].tolist()
|
||||
leading_toks_hash = hashlib.shake_128(str(leading_toks).encode()).hexdigest(4)
|
||||
|
||||
mtp_result_dict = {
|
||||
"token_ids":token_ids,
|
||||
"leading_toks_hash": leading_toks_hash,
|
||||
"num_fwd_evals": num_fwd_evals,
|
||||
"t_prefill": t_prefill,
|
||||
"t_gen": t_gen,
|
||||
"tokens_generated": tokens_generated,
|
||||
"toks_gend_incl_prefillplus1": toks_gend_incl_prefillplus1,
|
||||
"avg_effective_k": avg_effective_k,
|
||||
"effective_k_values": effective_k_values,
|
||||
"tps": tokens_generated / t_gen if t_gen > 0 else 0.0,
|
||||
}
|
||||
return mtp_result_dict
|
||||
|
||||
return input_ids if include_prompt else input_ids[:,prompt_size:]
|
||||
|
||||
@torch.inference_mode()
|
||||
def _prep_generate_args(
|
||||
self,
|
||||
model,
|
||||
input_ids: torch.Tensor,
|
||||
generation_config = None,
|
||||
model_kwargs: dict = None,
|
||||
):
|
||||
# Setup
|
||||
if model_kwargs is None:
|
||||
model_kwargs = {}
|
||||
|
||||
if generation_config is None:
|
||||
generation_config = model.generation_config
|
||||
model_kwargs["use_cache"] = True
|
||||
if "past_key_values" in model_kwargs:
|
||||
print(f"Before _get_initial_cache_position past_key_values cache length: {model_kwargs['past_key_values'].layers[0].get_seq_length()}", flush=True)
|
||||
model_kwargs = model._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
|
||||
if "past_key_values" in model_kwargs:
|
||||
print(f"After _get_initial_cache_position past_key_values cache length: {model_kwargs['past_key_values'].layers[0].get_seq_length()}", flush=True)
|
||||
return model_kwargs, generation_config
|
||||
|
||||
@torch.inference_mode()
|
||||
def _top1_confidence(
|
||||
self,
|
||||
logits: torch.Tensor = None
|
||||
):
|
||||
probs = F.softmax(logits, dim=-1) # LxV
|
||||
top1_idx = torch.argmax(probs, dim=-1) # Lx1
|
||||
top1_confs = probs[torch.arange(probs.shape[0], device=probs.device), top1_idx] # Lx1
|
||||
return top1_confs
|
||||
|
||||
@torch.inference_mode()
|
||||
def _extend_w_mask(
|
||||
self,
|
||||
input_ids=None,
|
||||
k_toks=None,
|
||||
mask_id=None,
|
||||
min_mask_id=None,
|
||||
max_mask_id=None
|
||||
):
|
||||
bsz, _ = input_ids.shape
|
||||
|
||||
if k_toks-1 > 0:
|
||||
if min_mask_id is not None:
|
||||
mask_tensor = torch.arange(min_mask_id, min_mask_id + k_toks - 1, dtype=torch.int64, device=input_ids.device).unsqueeze(0).expand(bsz, -1)
|
||||
# check that we didn't insert something larger than max_mask_id
|
||||
assert mask_tensor.max().item() <= max_mask_id, "Inserted mask ID exceeds specified max_mask_id."
|
||||
else:
|
||||
mask_tensor = torch.ones((bsz,k_toks-1),dtype=torch.int64, device=input_ids.device) * mask_id
|
||||
return torch.cat([input_ids, mask_tensor], dim=-1)
|
||||
|
||||
return input_ids
|
||||
|
||||
@torch.inference_mode()
|
||||
def _mtp_next_tokens(
|
||||
self,
|
||||
model,
|
||||
model_inputs,
|
||||
k_toks: int = 1,
|
||||
strategy: Optional[list] = None,
|
||||
) -> torch.Tensor:
|
||||
outputs = model(**model_inputs)
|
||||
# logits = outputs.logits[:, -k_toks:, :]
|
||||
logits = outputs.logits[0, -k_toks:, :]
|
||||
if strategy is None:
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
elif strategy[0] == "conf_adapt" or strategy[0] == "conf_adapt_sample@1":
|
||||
# we compute the position wise confidences using the _top1_confidence function
|
||||
top1_conf = self._top1_confidence(logits)
|
||||
# print(f"top1_conf: {top1_conf}", flush=True)
|
||||
# now we compute the position of the farthest token geq the threshold
|
||||
# but contiguously, so if we have [0.95, 0.92, 0.85, 0.97] and threshold 0.9
|
||||
# we want to get position 1 not 3, since position 2 is below the threshold
|
||||
# also being careful of situation like [0.85, 0.88, 0.95] where nothing meets the threshold
|
||||
# falling back to the first token in that case
|
||||
threshold = strategy[1]
|
||||
lt_thresh_mask = top1_conf < threshold
|
||||
# now we find the first case where the mask is true, and go back one position
|
||||
if torch.all(~lt_thresh_mask):
|
||||
# then all positions are above the threshold, we take the last position
|
||||
last_pos = k_toks - 1
|
||||
else:
|
||||
last_pos = torch.argmax(lt_thresh_mask.int()).item() - 1
|
||||
if last_pos < 0:
|
||||
last_pos = 0
|
||||
|
||||
# print(f"last_pos: {last_pos}", flush=True)
|
||||
|
||||
# then we slice the logits to only keep up to that position
|
||||
logits = logits[: last_pos + 1]
|
||||
|
||||
if last_pos == 0 and strategy[0] == "conf_adapt_sample@1":
|
||||
# if k is 1 this step, then draw from the distribution
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
# print(f"Sampling from probs at k={logits.size(0)}: {probs.shape}", flush=True)
|
||||
temperature = strategy[2]
|
||||
if 0.0 < temperature < 1.0:
|
||||
probs = probs.pow(1.0 / temperature)
|
||||
probs = probs / probs.sum(dim=-1, keepdim=True)
|
||||
else:
|
||||
print(f"Using temperature={temperature} has no effect.", flush=True)
|
||||
_next = torch.multinomial(probs, num_samples=1).squeeze(0)
|
||||
|
||||
else:
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
elif strategy[0] == "random":
|
||||
sampling_weights = strategy[1]
|
||||
# we sample k for this step according to the provided weights
|
||||
k_toks = int(
|
||||
torch.multinomial(torch.tensor(sampling_weights), num_samples=1).item()
|
||||
)
|
||||
logits = logits[: k_toks + 1]
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
else:
|
||||
raise ValueError(f"Unknown strategy: {strategy}")
|
||||
|
||||
_next = _next.unsqueeze(0) # add batch dim back
|
||||
return _next, outputs
|
||||
|
||||
|
||||
class LlamaForSequenceClassification(GenericForSequenceClassification, LlamaPreTrainedModel): ...
|
||||
|
||||
|
||||
class LlamaForQuestionAnswering(GenericForQuestionAnswering, LlamaPreTrainedModel):
|
||||
base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model`
|
||||
|
||||
|
||||
class LlamaForTokenClassification(GenericForTokenClassification, LlamaPreTrainedModel): ...
|
||||
|
||||
|
||||
#################################### HF registration ############################################################
|
||||
|
||||
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
|
||||
|
||||
LlamaConfig.register_for_auto_class()
|
||||
|
||||
LlamaForCausalLM.register_for_auto_class("AutoModel")
|
||||
LlamaForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
|
||||
__all__ = [
|
||||
"LlamaForCausalLM",
|
||||
"LlamaModel",
|
||||
"LlamaPreTrainedModel",
|
||||
"LlamaForSequenceClassification",
|
||||
"LlamaForQuestionAnswering",
|
||||
"LlamaForTokenClassification",
|
||||
]
|
||||
Reference in New Issue
Block a user