modeling_pebble.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

# Try to import the fast CUDA Mamba2.
try:
    from mamba_ssm import Mamba2
    HAS_MAMBA_SSM = True
except ImportError:
    HAS_MAMBA_SSM = False

from .configuration_pebble import PebbleConfig

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, 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.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, 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)
        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 PurePyTorchMamba2(nn.Module):
    """
    A pure PyTorch implementation of Mamba2 that exactly matches the 
    parameter names and math of mamba_ssm.Mamba2, allowing it to run on CPU.
    """
    def __init__(self, config):
        super().__init__()
        mamba_cfg = config.mamba2
        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)
        
        self.d_model = d_model
        self.d_state = d_state
        self.d_inner = expand * d_model
        self.headdim = headdim
        self.nheads = self.d_inner // headdim
        self.d_conv = d_conv
        
        self.in_proj = nn.Linear(d_model, 2 * self.d_inner + 2 * d_state + self.nheads, bias=False)
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
        self.norm = RMSNorm(self.d_inner, eps=config.rms_norm_eps)
        
        conv_in_channels = self.d_inner + 2 * d_state
        self.conv1d = nn.Conv1d(
            in_channels=conv_in_channels, 
            out_channels=conv_in_channels, 
            kernel_size=d_conv, 
            padding=d_conv-1, 
            groups=conv_in_channels, 
            bias=True
        )
        
        self.A_log = nn.Parameter(torch.zeros(self.nheads, dtype=torch.float32))
        self.D = nn.Parameter(torch.ones(self.nheads))
        self.dt_bias = nn.Parameter(torch.ones(self.nheads))
        
    def forward(self, x):
        B, T, C = x.shape
        xzbc = self.in_proj(x)
        
        x, z, B_p, C_p, dt = torch.split(
            xzbc, 
            [self.d_inner, self.d_inner, self.d_state, self.d_state, self.nheads], 
            dim=-1
        )
        
        xbc = torch.cat([x, B_p, C_p], dim=-1)
        xbc = xbc.transpose(1, 2)
        xbc = self.conv1d(xbc)[:, :, :T]
        xbc = xbc.transpose(1, 2)
        
        x, B_p, C_p = torch.split(
            xbc, 
            [self.d_inner, self.d_state, self.d_state], 
            dim=-1
        )
        
        x = F.silu(x)
        x = self.norm(x)
        z = self.norm(z)
        
        dt = F.softplus(dt + self.dt_bias)
        A = -torch.exp(self.A_log)
        
        # Reshape x for multi-head SSM
        x = x.view(B, T, self.nheads, self.headdim)
        
        B_p = B_p.view(B, T, 1, 1, self.d_state)
        C_p = C_p.view(B, T, 1, 1, self.d_state)
        A = A.view(1, self.nheads, 1, 1)
        
        h = torch.zeros(B, self.nheads, self.d_state, self.headdim, device=x.device)
        ys = []
        
        for t in range(T):
            dt_t = dt[:, t].view(B, self.nheads, 1, 1)
            dA = torch.exp(dt_t * A)
            
            dB = dt_t * B_p[:, t]
            x_t = x[:, t].unsqueeze(2)
            
            h = dA * h + (dB.transpose(-1, -2) * x_t)
            
            y = (h * C_p[:, t].transpose(-1, -2)).sum(dim=2)
            ys.append(y)
            
        y = torch.stack(ys, dim=1)
        y = y.view(B, T, self.d_inner)
        
        D = self.D.view(1, 1, self.nheads, 1)
        y = y + (x * D).view(B, T, self.d_inner)
        
        # z is already (B, T, d_inner), so it multiplies perfectly!
        out = y * z
        out = self.out_proj(out)
        return out


class MambaBlock(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.ln = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        mamba_cfg = config.mamba2
        
        if HAS_MAMBA_SSM and torch.cuda.is_available():
            # Fast CUDA path for NVIDIA GPUs
            self.use_hf_fallback = False
            self.mixer = Mamba2(
                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),
            )
        else:
            # Pure PyTorch fallback for Macs / CPUs / AMD GPUs
            self.use_hf_fallback = True
            self.mixer = PurePyTorchMamba2(config)

    def forward(self, x):
        return x + self.mixer(self.ln(x))

class PebbleForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = PebbleConfig
    supports_gradient_checkpointing = False
    _no_split_modules = ["MambaBlock", "AttentionBlock"]

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        
        self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
        
        self.blocks = nn.ModuleList([
            MambaBlock(config, layer_idx=i) if i % 4 < 3
            else AttentionBlock(config)
            for i in range(config.num_hidden_layers)
        ])
        
        self.lnf = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        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, 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:
            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,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
        }