modeling_cma.py

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
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"
    attribute_map = {
        "hidden_size": "d_model",
        "num_hidden_layers": "n_layers",
        "num_attention_heads": "n_heads",
        "num_key_value_heads": "n_kv_heads",
    }

    def __init__(
        self,
        vocab_size=4096,
        seq_len=1024,
        d_model=216,
        n_layers=10,
        n_heads=6,
        n_kv_heads=2,
        chunk=24,
        cma_heads=3,
        expand=2,
        cma_identity_prob=0.90,
        max_position_embeddings=None,
        n_positions=None,
        n_ctx=None,
        tie_word_embeddings=True,
        **kwargs,
    ):
        if min(vocab_size, seq_len, d_model, n_layers) <= 0:
            raise ValueError("Vocabulary, context, width, and depth must be positive.")
        if min(n_heads, n_kv_heads, chunk, cma_heads, expand) <= 0:
            raise ValueError("Attention, CMA, and expansion dimensions must be positive.")
        if d_model % n_heads or n_heads % n_kv_heads:
            raise ValueError(
                "d_model % n_heads and n_heads % n_kv_heads must be zero."
            )
        if (d_model // n_heads) % 2:
            raise ValueError("The token-attention head dimension must be even for RoPE.")
        if d_model % chunk or chunk % cma_heads:
            raise ValueError("CMA requires d_model % chunk == 0 and chunk % cma_heads == 0.")
        if d_model // chunk < 2:
            raise ValueError("CMA requires at least two channel chunks.")
        if not 0.0 < cma_identity_prob < 1.0:
            raise ValueError("cma_identity_prob must be between zero and one.")
        kwargs.setdefault("is_decoder", True)
        kwargs.setdefault("is_encoder_decoder", False)
        kwargs.setdefault("tie_word_embeddings", tie_word_embeddings)
        # CMA exports intentionally recompute the full visible context during
        # generation; this architecture does not expose a KV cache.
        kwargs.setdefault("use_cache", False)
        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
        self.head_dim = d_model // n_heads
        self.num_key_value_groups = n_heads // n_kv_heads
        self.is_decoder = True
        self.use_cache = False


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
        version = torch.__version__.split("+", 1)[0].split(".")
        self.native_gqa = tuple(map(int, version[:2])) >= (2, 5)
        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)
        if self.h != self.kv_h and not self.native_gqa:
            repeats = self.h // self.kv_h
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            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):
    # Pre-CMA: signed residual attention across channel chunks.

    def __init__(self, d, chunk=24, heads=3, 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 chunks.")
        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 _route(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)
        qk_in = xc + self.chunk_emb + global_state
        value_slots = self.wv(x).reshape(
            batch_tokens, self.n, self.h, self.hd, self.expand
        )
        key_in = value_slots.mean(dim=-1).reshape(batch_tokens, self.n, self.c)
        key_in = key_in + self.chunk_emb
        q = torch.einsum("bnc,nco->bno", qk_in, self.wqk[..., : self.c])
        k = torch.einsum("bnc,nco->bno", key_in, 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 @ k.transpose(-2, -1)) * scale.view(
            1, self.h, 1, 1
        ).to(q.dtype)
        logits = logits + self.bias.unsqueeze(0).to(q.dtype)
        attn = F.softmax(logits, dim=-1, dtype=torch.float32).to(v.dtype)
        routed = attn @ v
        route_signal = (
            torch.einsum(
                "bhnd,hd->bhn", q.float(), self.route_gate_weight.float()
            )
            + 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)
        return v, routed, route_coeff, logits, attn

    def forward(self, x):
        B, T, _ = x.shape
        v, routed, route_coeff, _, _ = self._route(x)
        y = v + route_coeff.unsqueeze(-1) * (routed - v)
        y = y.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]
        B, T, _ = x.shape
        v, routed, route_coeff, logits, attn = self._route(x)
        contribution = route_coeff.unsqueeze(-1) * (routed - v)
        mixed = v + contribution
        base = v.transpose(1, 2).reshape(B, T, self.d * self.expand)
        mixed = mixed.transpose(1, 2).reshape(B, T, self.d * self.expand)
        gate = F.silu(self.gate(x))
        base_output = self.o(base * gate)
        output = self.o(mixed * gate)
        flat_output = output.reshape(-1, self.d)
        token_effect = (
            flat_output.float().norm(dim=-1)
            / x.reshape(-1, self.d).float().norm(dim=-1).clamp_min(1e-12)
        )
        interaction_effect = (
            (output - base_output).reshape(-1, self.d).float().norm(dim=-1)
            / flat_output.float().norm(dim=-1).clamp_min(1e-12)
        )
        probs = attn.float().clamp_min(1e-9)
        entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
        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,
            "interaction_effect": interaction_effect,
            "probe_effect": interaction_effect.reshape(probe_count, -1).mean(dim=-1),
            "entropy": entropy.mean().item(),
            "diagonal_mass": probs.diagonal(dim1=-2, dim2=-1).mean().item(),
            "dominant_mass": probs.max(dim=-1).values.mean().item(),
            "head_similarity": head_similarity.item(),
            "gate_abs_mean": route_coeff.float().abs().mean().item(),
            "gate_saturation": (
                route_coeff.float().abs() > 0.95
            ).float().mean().item(),
            "contribution_ratio": (
                contribution.float().norm() / v.float().norm().clamp_min(1e-9)
            ).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
        compute_dtype = (
            torch.get_autocast_dtype(device_type)
            if torch.is_autocast_enabled(device_type)
            else x.dtype
        )
        x = x.to(dtype=compute_dtype)
        cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=compute_dtype)
        sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=compute_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"
    _no_split_modules = ["Block"]
    _tied_weights_keys = {"head.weight": "model.emb.weight"}
    # Transformers 4.x reads the expanded map directly; Transformers 5.x
    # replaces it during post_init(). Keeping both forms makes the same remote
    # code load cleanly across that boundary.
    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,
    ):
        return_dict = kwargs.pop(
            "return_dict", getattr(self.config, "use_return_dict", True)
        )
        if input_ids is None:
            raise ValueError("input_ids must be provided.")
        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,
            )
        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        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 :],
            "use_cache": False,
        }
        if attention_mask is not None:
            result["attention_mask"] = attention_mask[:, -self.config.seq_len :]
        return result