Update modeling_pebble.py

modeling_pebble.py

@@ -2,13 +2,14 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PreTrainedModel
-from .configuration_pebble import PebbleConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
 
 try:
     from mamba_ssm import Mamba2
 except ImportError:
-    Mamba2 = None
-    print("Warning: mamba-ssm not installed. Please install it to use PebbleLM.")
+    raise ImportError("mamba-ssm is required. pip install mamba-ssm causal-conv1d")
+
+from .configuration_pebble import PebbleConfig
 
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
@@ -23,101 +24,122 @@ class RMSNorm(nn.Module):
         return self.weight * xf.to(dt)
 
 class AttentionBlock(nn.Module):
-    def __init__(self, dim, n_heads, hidden, rope_theta=10000.0):
+    def __init__(self, config):
         super().__init__()
+        dim = config.hidden_size
+        n_heads = config.num_attention_heads
+        hidden = config.intermediate_size
         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)
+        self.ln1  = RMSNorm(dim, eps=config.rms_norm_eps)
+        self.ln2  = RMSNorm(dim, eps=config.rms_norm_eps)
+        self.rope_theta = config.attention.get("rope_theta", 10000.0)
 
     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)
+
+        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)
+            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)
+        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):
+    def __init__(self, config):
         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.ln = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        mamba_cfg = config.mamba2
         self.mixer = Mamba2(
-            d_model=dim,
-            d_state=d_state,
-            d_conv=d_conv,
-            expand=expand,
-            headdim=headdim,
-            use_mem_eff_path=True,
+            d_model=config.hidden_size,
+            d_state=mamba_cfg.get("d_state", 128),
+            d_conv=mamba_cfg.get("d_conv", 4),
+            expand=mamba_cfg.get("expand", 2),
+            headdim=mamba_cfg.get("headdim", 64),
+            use_mem_eff_path=mamba_cfg.get("use_mem_eff_path", True),
         )
 
     def forward(self, x):
         return x + self.mixer(self.ln(x))
 
-class PebbleLM(PreTrainedModel):
+class PebbleForCausalLM(PreTrainedModel):
     config_class = PebbleConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["MambaBlock", "AttentionBlock"]
 
     def __init__(self, config):
         super().__init__(config)
-        self.wte = nn.Embedding(config.vocab_size, config.d_model)
+        self.config = config
+        
+        self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
+        
+        # 3:1 Mamba:Attention ratio layout
         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)
+            MambaBlock(config) if i % 4 < 3
+            else AttentionBlock(config)
+            for i in range(config.num_hidden_layers)
         ])
-        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  # weight sharing
+        
+        self.lnf = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        # Tie weights
+        self.tie_weights()
+
+    def tie_weights(self):
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.wte.weight
 
-    def forward(self, input_ids=None, labels=None, targets=None, **kwargs):
+    def forward(self, input_ids=None, attention_mask=None, labels=None, past_key_values=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}
+            # Shift so that tokens < n predict n+1
+            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)
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values,
+        )
 
-# Register the model for AutoModel
-from transformers import AutoModelForCausalLM
-AutoModelForCausalLM.register(PebbleConfig, PebbleLM)
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+        }