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"]