Hugging Face's logo

KordAI
/
III-tinystories

Text Generation
Transformers
Safetensors
iii
causal-lm
language-model
small-language-model
tinystories
experimental
custom-architecture
custom_code
Model card Files
xet
 Community
III-tinystories / modeling_iii.py
Naphon's picture
Naphon
Fix RoPE problem
4984ba0 verified
Raw History Blame Contribute Delete
51.6 kB
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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 Optional
import torch
import torch.nn.functional as F
from torch import Tensor
from torch.nn import (
Module,
Parameter,
Linear,
SiLU,
Conv1d,
Embedding,
)
from transformers.cache_utils import (
Cache,
DynamicCache,
)
from transformers.modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
)
from transformers.modeling_utils import (
PreTrainedModel,
)
from transformers.generation import GenerationMixin
from transformers.pytorch_utils import (
ALL_LAYERNORM_LAYERS,
)
from .configuration_iii import IIIConfig
_CONFIG_FOR_DOC = "IIIConfig"
# ============================================================
# RMSNorm
# ============================================================
class IIIRMSNorm(Module):
def __init__(
self,
dim,
eps=1e-6,
):
super().__init__()
self.weight = Parameter(
torch.ones(dim)
)
self.epsilon = eps
def forward(
self,
tensor,
):
input_dtype = tensor.dtype
tensor = tensor.to(
torch.float32
)
variance = tensor.pow(2).mean(
-1,
keepdim=True,
)
tensor = tensor * torch.rsqrt(
variance + self.epsilon
)
return (
self.weight
* tensor.to(input_dtype)
)
ALL_LAYERNORM_LAYERS.append(
IIIRMSNorm
)
# ============================================================
# RoPE
# ============================================================
class IIIRoPE(Module):
def __init__(
self,
dim,
max_position_embeddings=100000,
theta=10000.0,
device=None,
):
super().__init__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.theta = theta
inv_freq = 1.0 / (
theta
** (
torch.arange(
0,
dim,
2,
device=device,
dtype=torch.float32,
)
/ dim
)
)
self.register_buffer(
"inv_freq",
inv_freq,
persistent=False,
)
# Do not build RoPE cache during initialization.
# This makes loading model weights independent of the cache.
self.register_buffer(
"cos_cached",
None,
persistent=False,
)
self.register_buffer(
"sin_cached",
None,
persistent=False,
)
self.max_seq_len_cached = 0
def _set_cos_sin_cache(
self,
seq_len,
device,
dtype,
):
self.max_seq_len_cached = seq_len
t = torch.arange(
seq_len,
device=device,
dtype=self.inv_freq.dtype,
)
freqs = torch.outer(
t,
self.inv_freq.to(device=t.device),
)
emb = torch.cat(
(freqs, freqs),
dim=-1,
)
self.cos_cached = emb.cos().to(dtype)
self.sin_cached = emb.sin().to(dtype)
def forward(
self,
x,
seq_len=None,
):
if seq_len is None:
seq_len = x.shape[-2]
# RoPE cache does not exist yet, or is too short.
if (
self.cos_cached is None
or self.sin_cached is None
or seq_len > self.max_seq_len_cached
):
self._set_cos_sin_cache(
seq_len=seq_len,
device=x.device,
dtype=x.dtype,
)
return (
self.cos_cached[:seq_len].to(
device=x.device,
dtype=x.dtype,
),
self.sin_cached[:seq_len].to(
device=x.device,
dtype=x.dtype,
),
)
def rotate_half(x):
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,
unsqueeze_dim=1,
):
cos = cos[
position_ids
].unsqueeze(
unsqueeze_dim
)
sin = sin[
position_ids
].unsqueeze(
unsqueeze_dim
)
q_embed = (
q * cos
+ rotate_half(q) * sin
)
k_embed = (
k * cos
+ rotate_half(k) * sin
)
return q_embed, k_embed
# ============================================================
# SwiGLU-style MLP
# ============================================================
class IIIMLP(Module):
def __init__(
self,
config,
dim=None,
intermediate_dim=None,
):
super().__init__()
self.config = config
self.dim = (
config.hidden_dim
if dim is None
else dim
)
self.intermediate_dim = (
config.intermediate_dim
if intermediate_dim is None
else intermediate_dim
)
self.gate = Linear(
self.dim,
self.intermediate_dim,
bias=False,
)
self.up = Linear(
self.dim,
self.intermediate_dim,
bias=False,
)
self.down = Linear(
self.intermediate_dim,
self.dim,
bias=False,
)
self.activation = SiLU()
def forward(self, x):
value = self.up(x)
gate = self.activation(
self.gate(x)
)
return self.down(
gate * value
)
# ============================================================
# Attention
# ============================================================
class IIIAttention(Module):
def __init__(
self,
config,
layer_id=None,
):
super().__init__()
self.config = config
self.layer_id = layer_id
self.hidden_dim = config.hidden_dim
self.lora_dim = config.lora_dim
self.num_heads = config.n_heads
if self.lora_dim % self.num_heads != 0:
raise ValueError(
f"lora_dim ({self.lora_dim}) "
f"must be divisible by "
f"n_heads ({self.num_heads})."
)
self.head_dim = self.lora_dim // self.num_heads
# ----------------------------------------------------
# QKV
# ----------------------------------------------------
self.qkv = Linear(
self.hidden_dim,
self.lora_dim * 3,
bias=config.attention_bias,
)
self.o = Linear(
self.lora_dim,
self.hidden_dim,
bias=config.attention_bias,
)
# ----------------------------------------------------
# Q/K normalization
# ----------------------------------------------------
self.q_norm = IIIRMSNorm(
self.head_dim,
eps=config.rms_norm_eps,
)
self.k_norm = IIIRMSNorm(
self.head_dim,
eps=config.rms_norm_eps,
)
# ----------------------------------------------------
# RoPE
# ----------------------------------------------------
self.rotary_emb = IIIRoPE(
dim=self.head_dim,
max_position_embeddings=config.max_position_embedding,
theta=config.rope_theta,
)
# ============================================================
# Causal mask
# ============================================================
def _make_causal_mask(
self,
q_len,
kv_len,
device,
):
"""
Returns a boolean causal mask of shape:
[1, 1, q_len, kv_len]
True = attention is allowed
False = attention is blocked
Works for both:
1. Full sequence
q_len = kv_len = S
2. KV-cache decoding
q_len < kv_len
Example:
q_len=1, kv_len=5
[1 1 1 1 1]
Example:
q_len=2, kv_len=5
[1 1 1 1 0]
[1 1 1 1 1]
"""
q_positions = torch.arange(
kv_len - q_len,
kv_len,
device=device,
)
k_positions = torch.arange(
kv_len,
device=device,
)
causal = (
k_positions.unsqueeze(0)
<= q_positions.unsqueeze(1)
)
return causal.unsqueeze(0).unsqueeze(0)
# ============================================================
# Prepare user attention mask
# ============================================================
def _prepare_attention_mask(
self,
attention_mask,
device,
):
"""
Converts attention_mask into an attention-compatible shape.
Supported input shapes:
[B, KV]
[B, Q, KV]
[B, 1, Q, KV]
Boolean masks:
True = allowed
False = blocked
Integer masks:
1 = allowed
0 = blocked
Float masks:
treated as additive attention bias.
"""
if attention_mask is None:
return None
# --------------------------------------------------------
# [B, KV]
# --------------------------------------------------------
if attention_mask.dim() == 2:
if attention_mask.dtype == torch.bool:
return attention_mask[
:, None, None, :
].to(device)
# Integer 0/1 attention masks
if not torch.is_floating_point(attention_mask):
return (
attention_mask != 0
)[:, None, None, :].to(device)
# Float mask.
# Hugging Face commonly uses 0 for valid tokens
# and a large negative value for masked tokens.
return attention_mask[
:, None, None, :
].to(device)
# --------------------------------------------------------
# [B, Q, KV]
# --------------------------------------------------------
if attention_mask.dim() == 3:
return attention_mask[
:, None, :, :
].to(device)
# --------------------------------------------------------
# [B, 1, Q, KV]
# --------------------------------------------------------
if attention_mask.dim() == 4:
return attention_mask.to(device)
raise ValueError(
"attention_mask must be "
"2D, 3D, or 4D."
)
# ============================================================
# Combine causal + attention mask
# ============================================================
def _combine_masks(
self,
attention_mask,
q_len,
kv_len,
device,
dtype,
):
"""
Combines the user's attention/padding mask with the
causal mask.
This is the important part.
We NEVER allow an attention_mask to replace causality.
Returns:
bool mask OR additive float mask
suitable for SDPA with is_causal=False.
"""
causal_mask = self._make_causal_mask(
q_len=q_len,
kv_len=kv_len,
device=device,
)
# --------------------------------------------------------
# No user mask
# --------------------------------------------------------
if attention_mask is None:
return causal_mask
mask = self._prepare_attention_mask(
attention_mask,
device,
)
# --------------------------------------------------------
# Boolean mask
# --------------------------------------------------------
if mask.dtype == torch.bool:
return mask & causal_mask
# --------------------------------------------------------
# Integer mask
# --------------------------------------------------------
if not torch.is_floating_point(mask):
mask = mask != 0
return mask & causal_mask
# --------------------------------------------------------
# Additive float mask
#
# Typical form:
#
# valid -> 0
# masked -> -inf / large negative
#
# We impose causality by setting all future positions
# to the minimum representable value.
# --------------------------------------------------------
min_value = torch.finfo(dtype).min
# Convert causal mask to additive form.
causal_bias = torch.zeros(
(1, 1, q_len, kv_len),
dtype=dtype,
device=device,
)
causal_bias = causal_bias.masked_fill(
~causal_mask,
min_value,
)
# Make sure the supplied mask has the correct dtype.
mask = mask.to(dtype)
return mask + causal_bias
# ============================================================
# Forward
# ============================================================
def forward(
self,
hidden_states,
attention_mask=None,
position_ids=None,
kv_cache: Optional[Cache] = None,
use_cache=False,
output_attentions=False,
):
batch_size, seq_len, _ = hidden_states.shape
# --------------------------------------------------------
# Validate cache usage
# --------------------------------------------------------
if (
use_cache
and kv_cache is not None
and self.layer_id is None
):
raise ValueError(
"layer_id must be set when "
"using KV cache."
)
# --------------------------------------------------------
# QKV
# --------------------------------------------------------
qkv = self.qkv(hidden_states)
q, k, v = qkv.chunk(
3,
dim=-1,
)
# --------------------------------------------------------
# Heads
# --------------------------------------------------------
q = q.view(
batch_size,
seq_len,
self.num_heads,
self.head_dim,
).transpose(1, 2)
k = k.view(
batch_size,
seq_len,
self.num_heads,
self.head_dim,
).transpose(1, 2)
v = v.view(
batch_size,
seq_len,
self.num_heads,
self.head_dim,
).transpose(1, 2)
# --------------------------------------------------------
# Q/K RMSNorm
# --------------------------------------------------------
q = self.q_norm(q)
k = self.k_norm(k)
# --------------------------------------------------------
# Cache length
# --------------------------------------------------------
if (
kv_cache is not None
and kv_cache.get_seq_length(
self.layer_id
) > 0
):
past_len = kv_cache.get_seq_length(
self.layer_id
)
else:
past_len = 0
# --------------------------------------------------------
# Position IDs
# --------------------------------------------------------
if position_ids is None:
position_ids = torch.arange(
past_len,
past_len + seq_len,
device=hidden_states.device,
)
position_ids = (
position_ids
.unsqueeze(0)
.expand(
batch_size,
-1,
)
)
max_position = (
int(position_ids.max().item())
+ 1
)
# --------------------------------------------------------
# RoPE
# --------------------------------------------------------
cos, sin = self.rotary_emb(
q,
seq_len=max_position,
)
q, k = apply_rotary_pos_emb(
q,
k,
cos,
sin,
position_ids,
)
# --------------------------------------------------------
# Update KV cache
# --------------------------------------------------------
if (
use_cache
and kv_cache is not None
):
k, v = kv_cache.update(
k,
v,
self.layer_id,
)
# --------------------------------------------------------
# Actual attention sizes
# --------------------------------------------------------
kv_len = k.shape[-2]
q_len = q.shape[-2]
# ========================================================
# Build combined mask
# ========================================================
mask = self._combine_masks(
attention_mask=attention_mask,
q_len=q_len,
kv_len=kv_len,
device=hidden_states.device,
dtype=q.dtype,
)
# We explicitly constructed the causal mask above,
# so don't also ask SDPA to construct another one.
#
# This is also necessary because SDPA does not allow
# simultaneously supplying an attn_mask and is_causal=True.
is_causal = False
# ========================================================
# Attention
# ========================================================
attention_weights = None
# --------------------------------------------------------
# Debug / visualization path
# --------------------------------------------------------
if output_attentions:
scale = self.head_dim ** -0.5
scores = torch.matmul(
q,
k.transpose(-2, -1),
) * scale
# ----------------------------------------------------
# Boolean mask
# ----------------------------------------------------
if mask.dtype == torch.bool:
scores = scores.masked_fill(
~mask,
torch.finfo(
scores.dtype
).min,
)
# ----------------------------------------------------
# Additive mask
# ----------------------------------------------------
else:
scores = scores + mask.to(
scores.dtype
)
attention_weights = torch.softmax(
scores,
dim=-1,
)
attn_output = torch.matmul(
attention_weights,
v,
)
# --------------------------------------------------------
# SDPA path
# --------------------------------------------------------
else:
attn_output = (
F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=mask,
dropout_p=0.0,
is_causal=is_causal,
)
)
# ========================================================
# Merge heads
# ========================================================
attn_output = (
attn_output
.transpose(1, 2)
.contiguous()
.view(
batch_size,
seq_len,
self.lora_dim,
)
)
# ========================================================
# Output
# ========================================================
attn_output = self.o(
attn_output
)
return (
attn_output,
kv_cache,
attention_weights,
)
# ============================================================
# Gated Causal CNN
# ============================================================
class GatedCausalCNN(Module):
"""
Gated causal dilated CNN.
value = W_v(x)
score = W_s(x)
gated = value * sigmoid(score)
output = CausalConv1D(gated)
No RoPE is applied here.
The convolution itself encodes relative local position
through its kernel offsets and dilation.
"""
def __init__(
self,
config,
layer_id,
):
super().__init__()
self.config = config
self.layer_id = layer_id
self.hidden_dim = (
config.hidden_dim
)
self.window_size = (
config.layer_window_sizes[
layer_id
]
)
self.dilation = (
config.layer_dilation_size[
layer_id
]
)
if self.window_size <= 0:
raise ValueError(
"GatedCausalCNN requires "
"window_size > 0."
)
if self.dilation <= 0:
raise ValueError(
"GatedCausalCNN requires "
"dilation > 0."
)
# Receptive field:
#
# 1 + (kernel_size - 1) * dilation
#
self.receptive_field = (
1
+ (
self.window_size - 1
) * self.dilation
)
# ----------------------------------------------------
# Gated projection
#
# One projection gives us:
#
# [value | score]
# ----------------------------------------------------
self.gate_proj = Linear(
self.hidden_dim,
self.hidden_dim * 2,
bias=False,
)
# ----------------------------------------------------
# Depthwise causal convolution
# ----------------------------------------------------
self.conv = Conv1d(
in_channels=self.hidden_dim,
out_channels=self.hidden_dim,
kernel_size=self.window_size,
dilation=self.dilation,
groups=self.hidden_dim,
bias=config.attention_bias,
)
# ----------------------------------------------------
# Output projection
# ----------------------------------------------------
self.out_proj = Linear(
self.hidden_dim,
self.hidden_dim,
bias=config.attention_bias,
)
def _gate(
self,
hidden_states,
):
"""
hidden_states:
[B, S, D]
returns:
[B, S, D]
"""
value, score = self.gate_proj(
hidden_states
).chunk(
2,
dim=-1,
)
return (
value
* torch.sigmoid(score)
)
def _causal_conv(
self,
x,
):
"""
x:
[B, S, D]
Returns:
[B, S, D]
"""
left_padding = (
self.receptive_field - 1
)
# [B, S, D]
# ↓
# [B, D, S]
x = x.transpose(
1,
2,
)
x = F.pad(
x,
(
left_padding,
0,
),
)
x = self.conv(x)
# [B, D, S]
# ↓
# [B, S, D]
return x.transpose(
1,
2,
)
def forward(
self,
hidden_states,
kv_cache: Optional[Cache] = None,
use_cache=False,
):
batch_size, seq_len, _ = (
hidden_states.shape
)
# ====================================================
# No cache
# ====================================================
if (
not use_cache
or kv_cache is None
):
gated = self._gate(
hidden_states
)
out = self._causal_conv(
gated
)
out = self.out_proj(out)
return (
out,
kv_cache,
)
# ====================================================
# Cache
# ====================================================
if self.layer_id is None:
raise ValueError(
"layer_id must be set when "
"using a cache."
)
past_len = (
kv_cache.get_seq_length(
self.layer_id
)
)
# ----------------------------------------------------
# Retrieve cached raw hidden states
#
# CNN caches store:
#
# [B, 1, S, D]
#
# in the key slot.
# ----------------------------------------------------
if past_len > 0:
layer_cache = (
kv_cache.layers[
self.layer_id
]
)
past_states = (
layer_cache.keys[:, 0]
)
# We only need the actual receptive
# field before the new token(s).
history_len = (
self.receptive_field - 1
)
past_states = (
past_states[
:, -history_len:, :
]
)
else:
past_states = (
hidden_states[:, :0, :]
)
# ----------------------------------------------------
# Concatenate history + new states
# ----------------------------------------------------
states = torch.cat(
(
past_states,
hidden_states,
),
dim=1,
)
# ----------------------------------------------------
# GATE FIRST
# ----------------------------------------------------
gated = self._gate(
states
)
# ----------------------------------------------------
# THEN CONV1D
# ----------------------------------------------------
out = self._causal_conv(
gated
)
# Only return the outputs belonging
# to the newly supplied tokens.
out = out[
:,
-seq_len:,
:,
]
out = self.out_proj(out)
# ----------------------------------------------------
# Cache RAW hidden states
#
# Do not cache the gated representation.
#
# The next token needs:
#
# value(history) *
# sigmoid(score(history))
#
# before convolution.
# ----------------------------------------------------
cache_keys = (
hidden_states.unsqueeze(1)
)
cache_values = (
hidden_states[
..., :1
].unsqueeze(1)
)
kv_cache.update(
cache_keys,
cache_values,
self.layer_id,
)
return (
out,
kv_cache,
)
# ============================================================
# III Decoder Layer
# ============================================================
class IIILayer(Module):
def __init__(
self,
config,
layer_id,
):
super().__init__()
self.config = config
self.layer_id = layer_id
self.window_size = (
config.layer_window_sizes[
layer_id
]
)
self.dilation = (
config.layer_dilation_size[
layer_id
]
)
# ----------------------------------------------------
# Pre-mixer RMSNorm
# ----------------------------------------------------
self.input_layernorm = (
IIIRMSNorm(
config.hidden_dim,
eps=config.rms_norm_eps,
)
)
# ----------------------------------------------------
# Mixer
# ----------------------------------------------------
if self.window_size == 0:
self.is_attention = True
self.mixer = IIIAttention(
config,
layer_id=layer_id,
)
else:
self.is_attention = False
self.mixer = GatedCausalCNN(
config,
layer_id=layer_id,
)
# ----------------------------------------------------
# Post-mixer RMSNorm
# ----------------------------------------------------
self.post_mixer_layernorm = (
IIIRMSNorm(
config.hidden_dim,
eps=config.rms_norm_eps,
)
)
# ----------------------------------------------------
# MLP
# ----------------------------------------------------
self.mlp = IIIMLP(config)
def forward(
self,
hidden_states,
attention_mask=None,
position_ids=None,
kv_cache=None,
use_cache=False,
output_attentions=False,
):
# ----------------------------------------------------
# Mixer
# ----------------------------------------------------
residual = hidden_states
hidden_states = (
self.input_layernorm(
hidden_states
)
)
if self.is_attention:
(
mixer_output,
kv_cache,
attention_weights,
) = self.mixer(
hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
kv_cache=kv_cache,
use_cache=use_cache,
output_attentions=(
output_attentions
),
)
else:
(
mixer_output,
kv_cache,
) = self.mixer(
hidden_states,
kv_cache=kv_cache,
use_cache=use_cache,
)
attention_weights = None
hidden_states = (
residual
+ mixer_output
)
# ----------------------------------------------------
# MLP
# ----------------------------------------------------
residual = hidden_states
hidden_states = (
self.post_mixer_layernorm(
hidden_states
)
)
hidden_states = self.mlp(
hidden_states
)
hidden_states = (
residual
+ hidden_states
)
return (
hidden_states,
kv_cache,
attention_weights,
)
# ============================================================
# Pretrained base
# ============================================================
class IIIPreTrainedModel(
PreTrainedModel
):
config_class = IIIConfig
base_model_prefix = "model"
_supports_cache_class = True
def _init_weights(
self,
module,
):
std = getattr(
self.config,
"initializer_range",
0.02,
)
if isinstance(
module,
Linear,
):
torch.nn.init.normal_(
module.weight,
mean=0.0,
std=std,
)
if module.bias is not None:
torch.nn.init.zeros_(
module.bias
)
elif isinstance(
module,
Embedding,
):
torch.nn.init.normal_(
module.weight,
mean=0.0,
std=std,
)
if (
module.padding_idx
is not None
):
torch.nn.init.zeros_(
module.weight[
module.padding_idx
]
)
elif isinstance(
module,
Conv1d,
):
torch.nn.init.normal_(
module.weight,
mean=0.0,
std=std,
)
if module.bias is not None:
torch.nn.init.zeros_(
module.bias
)
elif isinstance(
module,
IIIRMSNorm,
):
torch.nn.init.ones_(
module.weight
)
# ============================================================
# III Model
# ============================================================
class IIIModel(
IIIPreTrainedModel
):
def __init__(
self,
config,
):
super().__init__(
config
)
# ----------------------------------------------------
# Find an attention layer.
#
# This is important because CNN cache length can be
# local, while attention cache length represents the
# complete generated sequence.
# ----------------------------------------------------
attention_layers = [
i
for i, window in enumerate(
config.layer_window_sizes
)
if window == 0
]
if not attention_layers:
raise ValueError(
"IIIModel requires at least "
"one attention layer."
)
self.first_attention_layer = (
attention_layers[0]
)
# ----------------------------------------------------
# Token embeddings
# ----------------------------------------------------
self.embed_tokens = Embedding(
config.vocab_size,
config.hidden_dim,
padding_idx=config.pad_token_id,
)
# ----------------------------------------------------
# Decoder layers
# ----------------------------------------------------
self.layers = (
torch.nn.ModuleList(
[
IIILayer(
config,
layer_id=i,
)
for i in range(
config.n_layers
)
]
)
)
# ----------------------------------------------------
# Final norm
# ----------------------------------------------------
self.norm = IIIRMSNorm(
config.hidden_dim,
eps=config.rms_norm_eps,
)
self.gradient_checkpointing = (
False
)
self.post_init()
def get_input_embeddings(
self,
):
return self.embed_tokens
def set_input_embeddings(
self,
value,
):
self.embed_tokens = value
def _get_past_length(
self,
past_key_values,
):
if past_key_values is None:
return 0
# IMPORTANT:
#
# Layer 0 is a CNN and may conceptually only need
# a local history. Attention layers contain the
# complete sequence, so use the first attention
# layer to determine the real sequence position.
return past_key_values.get_seq_length(
self.first_attention_layer
)
def forward(
self,
input_ids=None,
attention_mask=None,
position_ids=None,
past_key_values=None,
inputs_embeds=None,
cache_position=None,
use_cache=None,
output_attentions=False,
output_hidden_states=False,
return_dict=None,
):
if (
input_ids is not None
and inputs_embeds is not None
):
raise ValueError(
"Specify either input_ids "
"or inputs_embeds, not both."
)
if (
input_ids is None
and inputs_embeds is None
):
raise ValueError(
"Either input_ids or "
"inputs_embeds must be provided."
)
if use_cache is None:
use_cache = (
self.config.use_cache
)
if return_dict is None:
return_dict = True
# ----------------------------------------------------
# Cache
# ----------------------------------------------------
if use_cache:
if past_key_values is None:
past_key_values = (
DynamicCache()
)
elif not isinstance(
past_key_values,
Cache,
):
raise TypeError(
"past_key_values must be "
"a Transformers Cache."
)
else:
past_key_values = None
# ----------------------------------------------------
# Inputs
# ----------------------------------------------------
if inputs_embeds is None:
hidden_states = (
self.embed_tokens(
input_ids
)
)
else:
hidden_states = (
inputs_embeds
)
batch_size, seq_len, _ = (
hidden_states.shape
)
# ----------------------------------------------------
# Positions
# ----------------------------------------------------
past_length = (
self._get_past_length(
past_key_values
)
)
if position_ids is None:
if cache_position is not None:
position_ids = (
cache_position
)
if position_ids.dim() == 1:
position_ids = (
position_ids
.unsqueeze(0)
.expand(
batch_size,
-1,
)
)
else:
position_ids = (
torch.arange(
past_length,
past_length
+ seq_len,
device=(
hidden_states.device
),
)
.unsqueeze(0)
.expand(
batch_size,
-1,
)
)
# ----------------------------------------------------
# Outputs
# ----------------------------------------------------
all_hidden_states = (
()
if output_hidden_states
else None
)
all_attentions = (
()
if output_attentions
else None
)
# ----------------------------------------------------
# Decoder
# ----------------------------------------------------
for layer in self.layers:
if output_hidden_states:
all_hidden_states += (
hidden_states,
)
(
hidden_states,
past_key_values,
attention_weights,
) = layer(
hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
kv_cache=past_key_values,
use_cache=use_cache,
output_attentions=(
output_attentions
),
)
if output_attentions:
all_attentions += (
attention_weights,
)
# ----------------------------------------------------
# Final norm
# ----------------------------------------------------
hidden_states = self.norm(
hidden_states
)
if output_hidden_states:
all_hidden_states += (
hidden_states,
)
# ----------------------------------------------------
# Tuple output
# ----------------------------------------------------
if not return_dict:
return tuple(
value
for value in (
hidden_states,
past_key_values,
all_hidden_states,
all_attentions,
)
if value is not None
)
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=(
past_key_values
if use_cache
else None
),
hidden_states=(
all_hidden_states
),
attentions=(
all_attentions
),
)
# ============================================================
# III Causal LM
# ============================================================
class IIIForCausalLM(
IIIPreTrainedModel,
GenerationMixin,
):
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
def __init__(
self,
config,
):
super().__init__(
config
)
self.model = IIIModel(
config
)
self.lm_head = Linear(
config.hidden_dim,
config.vocab_size,
bias=False,
)
self.post_init()
if config.tie_word_embedding:
self.lm_head.weight = (
self.model.embed_tokens.weight
)
# --------------------------------------------------------
# Embeddings
# --------------------------------------------------------
def get_input_embeddings(
self,
):
return self.model.embed_tokens
def set_input_embeddings(
self,
value,
):
self.model.embed_tokens = value
def get_output_embeddings(
self,
):
return self.lm_head
def set_output_embeddings(
self,
new_embeddings,
):
self.lm_head = new_embeddings
# --------------------------------------------------------
# Forward
# --------------------------------------------------------
def forward(
self,
input_ids=None,
attention_mask=None,
position_ids=None,
past_key_values=None,
inputs_embeds=None,
labels=None,
cache_position=None,
use_cache=None,
output_attentions=False,
output_hidden_states=False,
return_dict=None,
):
if return_dict is None:
return_dict = True
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
cache_position=cache_position,
use_cache=use_cache,
output_attentions=(
output_attentions
),
output_hidden_states=(
output_hidden_states
),
return_dict=True,
)
# ----------------------------------------------------
# LM logits
# ----------------------------------------------------
logits = self.lm_head(
outputs.last_hidden_state
)
# ----------------------------------------------------
# Causal LM loss
# ----------------------------------------------------
loss = None
if labels is not None:
shift_logits = (
logits[
...,
:-1,
:,
]
.contiguous()
)
shift_labels = (
labels[
...,
1:,
]
.contiguous()
)
loss = F.cross_entropy(
shift_logits.view(
-1,
shift_logits.size(-1),
),
shift_labels.view(
-1
),
)
# ----------------------------------------------------
# Tuple output
# ----------------------------------------------------
if not return_dict:
result = (
logits,
outputs.past_key_values,
outputs.hidden_states,
outputs.attentions,
)
if loss is not None:
result = (
loss,
*result,
)
return result
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=(
outputs.past_key_values
),
hidden_states=(
outputs.hidden_states
),
attentions=(
outputs.attentions
),
)
# --------------------------------------------------------
# Generation
# --------------------------------------------------------
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
attention_mask=None,
inputs_embeds=None,
cache_position=None,
**kwargs,
):
"""
Prepare only the unprocessed tokens for generation.
When a cache exists, the model should not recompute
hidden states for tokens already represented in the
cache.
"""
if (
past_key_values is not None
):
past_length = (
past_key_values.get_seq_length(
self.model.first_attention_layer
)
)
if (
input_ids is not None
and input_ids.shape[1]
> past_length
):
input_ids = input_ids[
:,
past_length:,
]
elif input_ids is not None:
# Usually generation reaches this path with
# one new token already supplied.
input_ids = input_ids[
:,
-1:,
]
if (
inputs_embeds is not None
):
inputs_embeds = (
inputs_embeds[
:,
-input_ids.shape[1]:,
:,
]
)
model_inputs = {}
if (
inputs_embeds is not None
and past_key_values is None
):
model_inputs[
"inputs_embeds"
] = inputs_embeds
if input_ids is not None:
model_inputs[
"input_ids"
] = input_ids
else:
model_inputs[
"input_ids"
] = input_ids
model_inputs[
"past_key_values"
] = past_key_values
model_inputs[
"attention_mask"
] = attention_mask
model_inputs[
"cache_position"
] = cache_position
model_inputs[
"use_cache"
] = kwargs.get(
"use_cache",
self.config.use_cache,
)
return model_inputs