Commit ·
123a250
1
Parent(s): 2334367
Update lance_ai_model.py
Browse files- lance_ai_model.py +1 -16
lance_ai_model.py
CHANGED
|
@@ -1,6 +1,3 @@
|
|
| 1 |
-
"""
|
| 2 |
-
LanceAI - RoPE, SwiGLU, RMSNorm, GQA architecture
|
| 3 |
-
"""
|
| 4 |
import math
|
| 5 |
import torch
|
| 6 |
import torch.nn as nn
|
|
@@ -55,7 +52,6 @@ class LanceAIConfig(PretrainedConfig):
|
|
| 55 |
self.bos_token_id = bos_token_id
|
| 56 |
self.eos_token_id = eos_token_id
|
| 57 |
|
| 58 |
-
|
| 59 |
class LanceAIRMSNorm(nn.Module):
|
| 60 |
def __init__(self, hidden_size, eps=1e-6):
|
| 61 |
super().__init__()
|
|
@@ -69,7 +65,6 @@ class LanceAIRMSNorm(nn.Module):
|
|
| 69 |
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 70 |
return self.weight * hidden_states.to(input_dtype)
|
| 71 |
|
| 72 |
-
|
| 73 |
class LanceAIRotaryEmbedding(nn.Module):
|
| 74 |
def __init__(self, config):
|
| 75 |
super().__init__()
|
|
@@ -89,13 +84,11 @@ class LanceAIRotaryEmbedding(nn.Module):
|
|
| 89 |
sin = emb.sin()
|
| 90 |
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 91 |
|
| 92 |
-
|
| 93 |
def rotate_half(x):
|
| 94 |
x1 = x[..., :x.shape[-1] // 2]
|
| 95 |
x2 = x[..., x.shape[-1] // 2:]
|
| 96 |
return torch.cat((-x2, x1), dim=-1)
|
| 97 |
|
| 98 |
-
|
| 99 |
def apply_rotary_pos_emb(q, k, cos, sin):
|
| 100 |
cos = cos.unsqueeze(1)
|
| 101 |
sin = sin.unsqueeze(1)
|
|
@@ -103,7 +96,6 @@ def apply_rotary_pos_emb(q, k, cos, sin):
|
|
| 103 |
k_embed = (k * cos) + (rotate_half(k) * sin)
|
| 104 |
return q_embed, k_embed
|
| 105 |
|
| 106 |
-
|
| 107 |
def repeat_kv(hidden_states, n_rep):
|
| 108 |
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 109 |
if n_rep == 1:
|
|
@@ -111,7 +103,6 @@ def repeat_kv(hidden_states, n_rep):
|
|
| 111 |
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 112 |
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 113 |
|
| 114 |
-
|
| 115 |
class LanceAIAttention(nn.Module):
|
| 116 |
def __init__(self, config, layer_idx):
|
| 117 |
super().__init__()
|
|
@@ -164,7 +155,6 @@ class LanceAIAttention(nn.Module):
|
|
| 164 |
|
| 165 |
return attn_output, past
|
| 166 |
|
| 167 |
-
|
| 168 |
class LanceAIMLP(nn.Module):
|
| 169 |
def __init__(self, config):
|
| 170 |
super().__init__()
|
|
@@ -177,7 +167,6 @@ class LanceAIMLP(nn.Module):
|
|
| 177 |
def forward(self, x):
|
| 178 |
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
| 179 |
|
| 180 |
-
|
| 181 |
class LanceAIDecoderLayer(nn.Module):
|
| 182 |
def __init__(self, config, layer_idx):
|
| 183 |
super().__init__()
|
|
@@ -206,7 +195,6 @@ class LanceAIDecoderLayer(nn.Module):
|
|
| 206 |
|
| 207 |
return hidden_states, past_kv
|
| 208 |
|
| 209 |
-
|
| 210 |
class LanceAIPreTrainedModel(PreTrainedModel):
|
| 211 |
config_class = LanceAIConfig
|
| 212 |
base_model_prefix = "model"
|
|
@@ -216,7 +204,6 @@ class LanceAIPreTrainedModel(PreTrainedModel):
|
|
| 216 |
_supports_flash_attn = True
|
| 217 |
_supports_sdpa = True
|
| 218 |
|
| 219 |
-
|
| 220 |
class LanceAIModel(LanceAIPreTrainedModel):
|
| 221 |
def __init__(self, config):
|
| 222 |
super().__init__(config)
|
|
@@ -312,7 +299,6 @@ class LanceAIModel(LanceAIPreTrainedModel):
|
|
| 312 |
mask = torch.triu(mask, diagonal=1 + past_len)
|
| 313 |
return mask[None, None, :, :]
|
| 314 |
|
| 315 |
-
|
| 316 |
class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
|
| 317 |
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
| 318 |
|
|
@@ -377,8 +363,7 @@ class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
|
|
| 377 |
reordered.append((layer_k.index_select(0, beam_idx), layer_v.index_select(0, beam_idx)))
|
| 378 |
return reordered
|
| 379 |
|
| 380 |
-
|
| 381 |
CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
|
| 382 |
MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)
|
| 383 |
LanceAIConfig.register_for_auto_class("AutoConfig")
|
| 384 |
-
LanceAI.register_for_auto_class("AutoModelForCausalLM")
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
import math
|
| 2 |
import torch
|
| 3 |
import torch.nn as nn
|
|
|
|
| 52 |
self.bos_token_id = bos_token_id
|
| 53 |
self.eos_token_id = eos_token_id
|
| 54 |
|
|
|
|
| 55 |
class LanceAIRMSNorm(nn.Module):
|
| 56 |
def __init__(self, hidden_size, eps=1e-6):
|
| 57 |
super().__init__()
|
|
|
|
| 65 |
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 66 |
return self.weight * hidden_states.to(input_dtype)
|
| 67 |
|
|
|
|
| 68 |
class LanceAIRotaryEmbedding(nn.Module):
|
| 69 |
def __init__(self, config):
|
| 70 |
super().__init__()
|
|
|
|
| 84 |
sin = emb.sin()
|
| 85 |
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 86 |
|
|
|
|
| 87 |
def rotate_half(x):
|
| 88 |
x1 = x[..., :x.shape[-1] // 2]
|
| 89 |
x2 = x[..., x.shape[-1] // 2:]
|
| 90 |
return torch.cat((-x2, x1), dim=-1)
|
| 91 |
|
|
|
|
| 92 |
def apply_rotary_pos_emb(q, k, cos, sin):
|
| 93 |
cos = cos.unsqueeze(1)
|
| 94 |
sin = sin.unsqueeze(1)
|
|
|
|
| 96 |
k_embed = (k * cos) + (rotate_half(k) * sin)
|
| 97 |
return q_embed, k_embed
|
| 98 |
|
|
|
|
| 99 |
def repeat_kv(hidden_states, n_rep):
|
| 100 |
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 101 |
if n_rep == 1:
|
|
|
|
| 103 |
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 104 |
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 105 |
|
|
|
|
| 106 |
class LanceAIAttention(nn.Module):
|
| 107 |
def __init__(self, config, layer_idx):
|
| 108 |
super().__init__()
|
|
|
|
| 155 |
|
| 156 |
return attn_output, past
|
| 157 |
|
|
|
|
| 158 |
class LanceAIMLP(nn.Module):
|
| 159 |
def __init__(self, config):
|
| 160 |
super().__init__()
|
|
|
|
| 167 |
def forward(self, x):
|
| 168 |
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
| 169 |
|
|
|
|
| 170 |
class LanceAIDecoderLayer(nn.Module):
|
| 171 |
def __init__(self, config, layer_idx):
|
| 172 |
super().__init__()
|
|
|
|
| 195 |
|
| 196 |
return hidden_states, past_kv
|
| 197 |
|
|
|
|
| 198 |
class LanceAIPreTrainedModel(PreTrainedModel):
|
| 199 |
config_class = LanceAIConfig
|
| 200 |
base_model_prefix = "model"
|
|
|
|
| 204 |
_supports_flash_attn = True
|
| 205 |
_supports_sdpa = True
|
| 206 |
|
|
|
|
| 207 |
class LanceAIModel(LanceAIPreTrainedModel):
|
| 208 |
def __init__(self, config):
|
| 209 |
super().__init__(config)
|
|
|
|
| 299 |
mask = torch.triu(mask, diagonal=1 + past_len)
|
| 300 |
return mask[None, None, :, :]
|
| 301 |
|
|
|
|
| 302 |
class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
|
| 303 |
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
| 304 |
|
|
|
|
| 363 |
reordered.append((layer_k.index_select(0, beam_idx), layer_v.index_select(0, beam_idx)))
|
| 364 |
return reordered
|
| 365 |
|
|
|
|
| 366 |
CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
|
| 367 |
MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)
|
| 368 |
LanceAIConfig.register_for_auto_class("AutoConfig")
|
| 369 |
+
LanceAI.register_for_auto_class("AutoModelForCausalLM")
|