| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import GenerationMixin, PreTrainedModel |
| from transformers.cache_utils import Cache, DynamicCache |
| from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast |
|
|
| from .configuration_smallm import SmallmConfig |
|
|
|
|
| def rms_norm(x): |
| return F.rms_norm(x, (x.size(-1),)) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, head_dim, max_seq_len, base): |
| super().__init__() |
| self.head_dim = head_dim |
| self.max_seq_len = max_seq_len |
| self.base = base |
| self.tables = None |
|
|
| def angles(self, device): |
| if self.tables is None or self.tables[0].device != device: |
| steps = torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) |
| freqs = self.base ** (-steps / self.head_dim) |
| positions = torch.arange(self.max_seq_len, dtype=torch.float32, device=device) |
| angles = torch.outer(positions, freqs) |
| self.tables = (angles.cos(), angles.sin()) |
| return self.tables |
|
|
| def forward(self, x, offset=0): |
| length = x.size(1) |
| cos_table, sin_table = self.angles(x.device) |
| cos = cos_table[offset:offset + length].view(1, length, 1, -1) |
| sin = sin_table[offset:offset + length].view(1, length, 1, -1) |
| left, right = x.float().chunk(2, dim=-1) |
| rotated = torch.cat([left * cos - right * sin, left * sin + right * cos], dim=-1) |
| return rotated.type_as(x) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, config, layer_idx): |
| super().__init__() |
| self.layer_idx = layer_idx |
| self.n_head = config.n_head |
| self.head_dim = config.dim // config.n_head |
| self.qkv = nn.Parameter(torch.empty(3, config.dim, config.dim)) |
| self.out = nn.Parameter(torch.empty(config.dim, config.dim)) |
| self.value_lambda = nn.Parameter(torch.tensor(0.5)) |
|
|
| def forward(self, x, first_value, rotary, attention_mask, past_key_values, offset): |
| batch, length, dim = x.shape |
| projected = F.linear(x, self.qkv.flatten(end_dim=1)) |
| q, k, v = projected.view(batch, length, 3, self.n_head, self.head_dim).unbind(dim=2) |
| q, k = rms_norm(q), rms_norm(k) |
| q, k = rotary(q, offset), rotary(k, offset) |
| if first_value is None: |
| first_value = v |
| v = torch.lerp(v, first_value, self.value_lambda) |
| q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
| if past_key_values is not None: |
| k, v = past_key_values.update(k, v, self.layer_idx) |
| attended = F.scaled_dot_product_attention( |
| q, k, v, |
| attn_mask=attention_mask, |
| is_causal=attention_mask is None and q.size(2) == k.size(2), |
| ) |
| attended = attended.transpose(1, 2).reshape(batch, length, dim) |
| return F.linear(attended, self.out), first_value |
|
|
|
|
| class Mlp(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.relu2_clamp = config.relu2_clamp |
| self.up = nn.Parameter(torch.empty(config.mlp_mult * config.dim, config.dim)) |
| self.down = nn.Parameter(torch.empty(config.dim, config.mlp_mult * config.dim)) |
|
|
| def forward(self, x): |
| hidden = F.relu(F.linear(x, self.up)) |
| return F.linear(hidden.clamp(max=self.relu2_clamp).square(), self.down) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config, layer_idx): |
| super().__init__() |
| self.attention = Attention(config, layer_idx) |
| self.mlp = Mlp(config) |
| self.residual_lambda = nn.Parameter(torch.tensor([1.0, 0.0])) |
|
|
| def forward(self, x, first_value, embedded, rotary, attention_mask, past_key_values, offset): |
| x = self.residual_lambda[0] * x + self.residual_lambda[1] * embedded |
| attended, first_value = self.attention( |
| rms_norm(x), first_value, rotary, attention_mask, past_key_values, offset |
| ) |
| x = x + attended |
| x = x + self.mlp(rms_norm(x)) |
| return x, first_value |
|
|
|
|
| class SmallmPreTrainedModel(PreTrainedModel): |
| config_class = SmallmConfig |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = False |
| _no_split_modules = ["Block"] |
|
|
| def _init_weights(self, module): |
| bound = self.config.dim ** -0.5 |
| if isinstance(module, Attention): |
| module.qkv.data.uniform_(-bound, bound) |
| module.out.data.zero_() |
| module.value_lambda.data.fill_(0.5) |
| elif isinstance(module, Mlp): |
| module.up.data.uniform_(-bound, bound) |
| module.down.data.zero_() |
| elif isinstance(module, Block): |
| module.residual_lambda.data.copy_(torch.tensor([1.0, 0.0])) |
| elif isinstance(module, nn.Embedding): |
| module.weight.data.uniform_(-0.5 * bound, 0.5 * bound) |
| elif isinstance(module, nn.Linear): |
| module.weight.data.zero_() |
|
|
|
|
| class SmallmModel(SmallmPreTrainedModel): |
| def __init__(self, config): |
| super().__init__(config) |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.dim) |
| self.rotary = RotaryEmbedding(config.dim // config.n_head, config.max_seq_len, |
| config.rope_theta) |
| self.layers = nn.ModuleList([Block(config, i) for i in range(config.n_layer)]) |
| self.n_skip = config.n_layer // 2 |
| self.skip_weights = nn.Parameter(torch.ones(self.n_skip)) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.embed_tokens = value |
|
|
| def forward(self, input_ids=None, attention_mask=None, past_key_values=None, |
| inputs_embeds=None, use_cache=None, **kwargs): |
| if inputs_embeds is None: |
| inputs_embeds = self.embed_tokens(input_ids) |
| use_cache = self.config.use_cache if use_cache is None else use_cache |
| if use_cache and past_key_values is None: |
| past_key_values = DynamicCache() |
| length = inputs_embeds.size(1) |
| offset = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| mask = self.build_attention_mask(attention_mask, length, offset + length, |
| inputs_embeds.device) |
|
|
| x = rms_norm(inputs_embeds) |
| embedded, first_value = x, None |
| skipped = [] |
| for i, layer in enumerate(self.layers): |
| if i >= self.n_skip: |
| x = x + self.skip_weights[i - self.n_skip] * skipped.pop() |
| x, first_value = layer(x, first_value, embedded, self.rotary, mask, |
| past_key_values if use_cache else None, offset) |
| if i < self.n_skip: |
| skipped.append(x) |
| return BaseModelOutputWithPast( |
| last_hidden_state=rms_norm(x), |
| past_key_values=past_key_values if use_cache else None, |
| ) |
|
|
| @staticmethod |
| def build_attention_mask(attention_mask, query_length, key_length, device): |
| if attention_mask is not None and attention_mask.dim() == 4: |
| return attention_mask |
| padded = attention_mask is not None and not bool(attention_mask.all()) |
| if not padded and query_length == key_length: |
| return None |
| keys = torch.arange(key_length, device=device) |
| queries = torch.arange(key_length - query_length, key_length, device=device) |
| allowed = (keys[None, :] <= queries[:, None])[None] |
| if attention_mask is not None: |
| allowed = allowed & attention_mask[:, None, :].bool() |
| allowed = allowed | (queries[None, :, None] == keys[None, None, :]) |
| return allowed.unsqueeze(1) |
|
|
|
|
| class SmallmForCausalLM(SmallmPreTrainedModel, GenerationMixin): |
| _tied_weights_keys = [] |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.model = SmallmModel(config) |
| self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| 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 forward(self, input_ids=None, attention_mask=None, past_key_values=None, |
| inputs_embeds=None, labels=None, use_cache=None, logits_to_keep=0, **kwargs): |
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| past_key_values=past_key_values, |
| inputs_embeds=inputs_embeds, |
| use_cache=use_cache, |
| ) |
| hidden = outputs.last_hidden_state |
| if isinstance(logits_to_keep, int) and logits_to_keep > 0: |
| hidden = hidden[:, -logits_to_keep:] |
| elif isinstance(logits_to_keep, torch.Tensor): |
| hidden = hidden[:, logits_to_keep] |
| logits = self.lm_head(hidden) |
| logits = self.config.softcap * torch.tanh(logits.float() / self.config.softcap) |
|
|
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits[:, :-1].reshape(-1, logits.size(-1)), |
| labels[:, 1:].reshape(-1).to(logits.device), |
| ) |
| return CausalLMOutputWithPast( |
| loss=loss, |
| logits=logits, |
| past_key_values=outputs.past_key_values, |
| ) |
|
|
|
|
| __all__ = ["SmallmConfig", "SmallmModel", "SmallmForCausalLM", "SmallmPreTrainedModel"] |
|
|
