| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel |
| from .configuration_pebble import PebbleConfig |
|
|
| try: |
| from mamba_ssm import Mamba2 |
| except ImportError: |
| Mamba2 = None |
| print("Warning: mamba-ssm not installed. Please install it to use PebbleLM.") |
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| dt = x.dtype |
| xf = x.float() |
| xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) |
| return self.weight * xf.to(dt) |
|
|
| class AttentionBlock(nn.Module): |
| def __init__(self, dim, n_heads, hidden, rope_theta=10000.0): |
| super().__init__() |
| assert dim % n_heads == 0 |
| self.nh, self.hd = n_heads, dim // n_heads |
| self.rope_theta = rope_theta |
| self.wqkv = nn.Linear(dim, 3 * dim, bias=False) |
| self.wo = nn.Linear(dim, dim, bias=False) |
| self.fc1 = nn.Linear(dim, hidden, bias=False) |
| self.fc2 = nn.Linear(hidden, dim, bias=False) |
| self.ln1 = RMSNorm(dim) |
| self.ln2 = RMSNorm(dim) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| h = self.ln1(x) |
| qkv = self.wqkv(h).view(B, T, 3, self.nh, self.hd).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
|
|
| half = self.hd // 2 |
| invf = 1.0 / (self.rope_theta ** ( |
| torch.arange(0, half, device=x.device, dtype=torch.float32) * 2.0 / self.hd)) |
| ang = torch.outer(torch.arange(T, device=x.device, dtype=torch.float32), invf) |
| cos, sin = ang.cos()[None, None], ang.sin()[None, None] |
|
|
| q1, q2 = q.float()[..., :half], q.float()[..., half:] |
| k1, k2 = k.float()[..., :half], k.float()[..., half:] |
| |
| q = torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1).to(v.dtype) |
| k = torch.cat([k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1).to(v.dtype) |
|
|
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| y = y.transpose(1, 2).reshape(B, T, C) |
| x = x + self.wo(y) |
| x = x + self.fc2(F.gelu(self.fc1(self.ln2(x)))) |
| return x |
|
|
| class MambaBlock(nn.Module): |
| def __init__(self, dim, d_state=128, d_conv=4, expand=2, headdim=64): |
| super().__init__() |
| if Mamba2 is None: |
| raise ImportError("mamba-ssm is not installed. Please install via `pip install mamba-ssm`") |
| self.ln = RMSNorm(dim) |
| self.mixer = Mamba2( |
| d_model=dim, |
| d_state=d_state, |
| d_conv=d_conv, |
| expand=expand, |
| headdim=headdim, |
| use_mem_eff_path=True, |
| ) |
|
|
| def forward(self, x): |
| return x + self.mixer(self.ln(x)) |
|
|
| class PebbleLM(PreTrainedModel): |
| config_class = PebbleConfig |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = True |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.wte = nn.Embedding(config.vocab_size, config.d_model) |
| self.blocks = nn.ModuleList([ |
| MambaBlock( |
| config.d_model, |
| config.mamba_d_state, |
| config.mamba_d_conv, |
| config.mamba_expand, |
| config.mamba_headdim |
| ) if i % 4 < 3 else AttentionBlock( |
| config.d_model, |
| config.n_heads, |
| config.att_hidden, |
| config.rope_theta |
| ) |
| for i in range(config.n_blocks) |
| ]) |
| self.lnf = RMSNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.lm_head.weight = self.wte.weight |
|
|
| def forward(self, input_ids=None, labels=None, targets=None, **kwargs): |
| x = self.wte(input_ids) |
| for blk in self.blocks: |
| x = blk(x) |
| logits = self.lm_head(self.lnf(x)) |
| |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.reshape(-1)) |
| elif targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1)) |
| |
| return {"logits": logits, "loss": loss} |
|
|
| |
| from transformers import AutoModelForCausalLM |
| AutoModelForCausalLM.register(PebbleConfig, PebbleLM) |
