modeling_cma.py

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
# CMA is text-only. Disable Transformers' optional vision backend before it
# imports PreTrainedModel: some cloud images expose a system torchvision built
# against a different PyTorch, which otherwise crashes on torchvision::nms.
import transformers as _transformers
import transformers.utils as _transformers_utils
from transformers.utils import import_utils as _transformers_import_utils


def _cma_torchvision_unavailable():
    return False


_transformers_import_utils._torchvision_available = False
_transformers_import_utils.is_torchvision_available = _cma_torchvision_unavailable
_transformers_import_utils.is_torchvision_v2_available = _cma_torchvision_unavailable
_transformers_utils.is_torchvision_available = _cma_torchvision_unavailable
_transformers_utils.is_torchvision_v2_available = _cma_torchvision_unavailable
_transformers.is_torchvision_available = _cma_torchvision_unavailable
_transformers.is_torchvision_v2_available = _cma_torchvision_unavailable
try:
    from transformers import GenerationMixin
except ImportError:
    from transformers.generation import GenerationMixin
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class CMAConfig(PretrainedConfig):
    model_type = "cma"

    def __init__(
        self,
        vocab_size=260,
        seq_len=2048,
        d_model=128,
        n_layers=6,
        n_heads=4,
        n_kv_heads=2,
        chunk=16,
        cma_heads=2,
        expand=2,
        cma_identity_prob=0.90,
        max_position_embeddings=None,
        n_positions=None,
        n_ctx=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.max_position_embeddings = max_position_embeddings or seq_len
        self.n_positions = n_positions or self.max_position_embeddings
        self.n_ctx = n_ctx or self.max_position_embeddings
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.chunk = chunk
        self.cma_heads = cma_heads
        self.expand = expand
        self.cma_identity_prob = cma_identity_prob


class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return F.rms_norm(
            x, (x.shape[-1],), self.w.to(dtype=x.dtype), eps=1e-6
        )


def rope_cache(seq_len, head_dim, device, base=10000.0):
    inv = 1.0 / (
        base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
    )
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv)
    return torch.cos(freqs), torch.sin(freqs)


def apply_rope(x, cos, sin):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


class TokenAttention(nn.Module):
    def __init__(self, d, n_heads, n_kv_heads):
        super().__init__()
        if d % n_heads or n_heads % n_kv_heads:
            raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero.")
        self.h, self.kv_h, self.hd = n_heads, n_kv_heads, d // n_heads
        self.q = nn.Linear(d, d, bias=False)
        self.k = nn.Linear(d, n_kv_heads * self.hd, bias=False)
        self.v = nn.Linear(d, n_kv_heads * self.hd, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)

    def forward(self, x, cos, sin):
        B, T, d = x.shape
        q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.k(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
        v = self.v(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
        q, k = self.qn(q), self.kn(k)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        # Native GQA keeps compact K/V heads and lets SDPA select a fused CUDA
        # implementation without materializing repeated K/V activations.
        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            is_causal=True,
            enable_gqa=self.h != self.kv_h,
        )
        return self.o(y.transpose(1, 2).reshape(B, T, d))


class CMA(nn.Module):
    def __init__(self, d, chunk=16, heads=2, expand=2, identity_prob=0.90):
        super().__init__()
        if min(d, chunk, heads, expand) <= 0:
            raise ValueError("CMA dimensions and expansion must be positive.")
        if d % chunk or chunk % heads:
            raise ValueError("CMA requires d % chunk == 0 and chunk % heads == 0.")
        if d // chunk < 2:
            raise ValueError("CMA requires at least two channel slots.")
        if not 0.0 < identity_prob < 1.0:
            raise ValueError("cma_identity_prob must be between zero and one.")

        self.d, self.n, self.c, self.h = d, d // chunk, chunk, heads
        self.hd, self.expand = chunk // heads, expand
        self.chunk_emb = nn.Parameter(torch.randn(self.n, chunk) * 0.02)
        self.wqk = nn.Parameter(torch.randn(self.n, chunk, 2 * chunk) * 0.02)
        self.global_proj = nn.Linear(d, chunk, bias=False)
        self.wv = nn.Linear(d, d * expand, bias=False)
        self.bias = nn.Parameter(torch.zeros(heads, self.n, self.n))
        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)
        self.logit_scale = nn.Parameter(torch.zeros(heads))
        self.layer_gain = nn.Parameter(torch.zeros(heads))
        self.route_gate_weight = nn.Parameter(torch.randn(heads, self.hd) * 0.02)
        self.route_gate_bias = nn.Parameter(torch.zeros(heads))
        self.gate = nn.Linear(d, d * expand, bias=False)
        self.o = nn.Linear(d * expand, d, bias=False)
        nn.init.zeros_(self.o.weight)

        diagonal_bias = math.log(
            (self.n - 1) * identity_prob / (1.0 - identity_prob)
        )
        with torch.no_grad():
            self.bias.add_(torch.eye(self.n) * diagonal_bias)

    def _routing(self, x):
        batch_tokens = x.numel() // self.d
        xc = x.reshape(batch_tokens, self.n, self.c)
        global_state = self.global_proj(x).reshape(batch_tokens, 1, self.c)
        q_input = xc + self.chunk_emb + global_state
        value_slots = self.wv(x).reshape(
            batch_tokens, self.n, self.h, self.hd, self.expand
        )
        key_input = value_slots.mean(dim=-1).reshape(batch_tokens, self.n, self.c)
        key_input = key_input + self.chunk_emb

        q = torch.einsum("bnc,nco->bno", q_input, self.wqk[..., : self.c])
        k = torch.einsum("bnc,nco->bno", key_input, self.wqk[..., self.c :])
        v = value_slots.reshape(
            batch_tokens, self.n, self.h, self.hd * self.expand
        ).transpose(1, 2)
        q = self.qn(q.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2)
        k = self.kn(k.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2)
        q = F.normalize(q.float(), dim=-1).to(v.dtype)
        k = F.normalize(k.float(), dim=-1).to(v.dtype)

        scale = self.logit_scale.clamp(max=math.log(100.0)).exp()
        logits = (q * scale.view(1, self.h, 1, 1).to(q.dtype)) @ k.transpose(-2, -1)
        logits = logits + self.bias.unsqueeze(0).to(logits.dtype)
        attn = F.softmax(logits, dim=-1, dtype=torch.float32).to(v.dtype)
        routed = attn @ v
        route_signal = (
            (q.float() * self.route_gate_weight.float().view(1, self.h, 1, self.hd)).sum(-1)
            + self.route_gate_bias.float().view(1, self.h, 1)
            + self.layer_gain.float().view(1, self.h, 1)
        )
        route_coeff = torch.tanh(route_signal).to(v.dtype)
        contribution = route_coeff.unsqueeze(-1) * (routed - v)
        return v, attn, route_coeff, contribution, logits

    def forward(self, x):
        B, T, _ = x.shape
        v, _, _, contribution, _ = self._routing(x)
        y = (v + contribution).transpose(1, 2).reshape(B, T, self.d * self.expand)
        return self.o(y * F.silu(self.gate(x)))

    @torch.no_grad()
    def diagnostic_stats(self, x):
        probe_count = x.shape[0] if x.ndim > 2 else 1
        v, attn, route_coeff, contribution, logits = self._routing(x)
        B, T, _ = x.shape
        base = v.transpose(1, 2).reshape(B, T, self.d * self.expand)
        routed = (v + contribution).transpose(1, 2).reshape(
            B, T, self.d * self.expand
        )
        gate = F.silu(self.gate(x))
        base_output = self.o(base * gate)
        routed_output = self.o(routed * gate)
        output_delta = (routed_output - base_output).reshape(-1, self.d)
        token_effect = (
            output_delta.float().norm(dim=-1)
            / routed_output.reshape(-1, self.d).float().norm(dim=-1).clamp_min(1e-12)
        )
        probe_effect = token_effect.reshape(probe_count, -1).mean(dim=-1)
        probs = attn.float().clamp_min(1e-9)
        entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
        diagonal_mass = probs.diagonal(dim1=-2, dim2=-1).mean()
        if self.h > 1:
            flattened = probs.transpose(0, 1).reshape(self.h, -1)
            normalized = F.normalize(flattened, dim=-1)
            similarity = normalized @ normalized.mT
            head_similarity = (
                similarity.sum() - similarity.diagonal().sum()
            ) / (self.h * (self.h - 1))
        else:
            head_similarity = probs.new_zeros(())
        return {
            "token_effect": token_effect,
            "probe_effect": probe_effect,
            "entropy": entropy.mean().item(),
            "diagonal_mass": diagonal_mass.item(),
            "head_similarity": head_similarity.item(),
            "gate_abs": route_coeff.float().abs().mean().item(),
            "gate_saturation": (route_coeff.float().abs() > 0.95).float().mean().item(),
            "logit_max": logits.float().abs().max().item(),
        }


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config.d_model
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
        self.attn = TokenAttention(d, config.n_heads, config.n_kv_heads)
        self.mix = CMA(
            d,
            config.chunk,
            config.cma_heads,
            config.expand,
            config.cma_identity_prob,
        )

    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        x = x + self.mix(self.n2(x))
        return x


class CMAModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config.d_model
        self.config = config
        self.emb = nn.Embedding(config.vocab_size, d)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
        self.norm = RMSNorm(d)
        hd = d // config.n_heads
        cos, sin = rope_cache(config.seq_len, hd, "cpu")
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)

    def forward_hidden(self, idx):
        if idx.size(1) > self.config.seq_len:
            idx = idx[:, -self.config.seq_len :]
        x = self.emb(idx)
        device_type = x.device.type
        rope_dtype = (
            torch.get_autocast_dtype(device_type)
            if torch.is_autocast_enabled(device_type)
            else x.dtype
        )
        cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
        sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
        for b in self.blocks:
            x = b(x, cos, sin)
        return self.norm(x)


class CMAForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = CMAConfig
    base_model_prefix = "model"
    _tied_weights_keys = {"head.weight": "model.emb.weight"}
    all_tied_weights_keys = {"head.weight": "model.emb.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.model = CMAModel(config)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # Modern Transformers creates loader metadata and performs configured
        # tying in post_init(); omitting it leaves all_tied_weights_keys absent.
        self.post_init()
        self.head.weight = self.model.emb.weight

    def get_input_embeddings(self):
        return self.model.emb

    def set_input_embeddings(self, value):
        self.model.emb = value

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, new_embeddings):
        self.head = new_embeddings

    def raw_logits(self, idx):
        return self.head(self.model.forward_hidden(idx))

    def logits(self, idx):
        return self.raw_logits(idx)

    def _masked_logits(self, input_ids, attention_mask):
        if attention_mask is None or bool(attention_mask.all()):
            return self.logits(input_ids)

        B, T = input_ids.shape
        out = None
        for i in range(B):
            keep = attention_mask[i].bool().nonzero(as_tuple=False).flatten()
            if keep.numel() == 0:
                keep = torch.tensor([T - 1], device=input_ids.device)
            trimmed = input_ids[i, keep].unsqueeze(0)
            logits_i = self.logits(trimmed)
            if out is None:
                out = logits_i.new_zeros(B, T, logits_i.size(-1))
            out[i, keep, :] = logits_i[0, -keep.numel() :, :]
        return out

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        use_cache=False,
        past_key_values=None,
        **kwargs,
    ):
        if input_ids.size(1) > self.config.seq_len:
            input_ids = input_ids[:, -self.config.seq_len :]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -self.config.seq_len :]
            if labels is not None:
                labels = labels[:, -self.config.seq_len :]
        logits = self._masked_logits(input_ids, attention_mask)
        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)).float(),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        attention_mask = kwargs.get("attention_mask")
        result = {"input_ids": input_ids[:, -self.config.seq_len :]}
        if attention_mask is not None:
            result["attention_mask"] = attention_mask[:, -self.config.seq_len :]
        return result