modeling_modernbert_small.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput

from babylm.model.attention import build_global_attention_mask, build_local_attention_mask, resolve_attention_types
from babylm.model.configuration_modernbert_small import ModernBertSmallConfig
from babylm.model.embeddings import get_embedding_class
from babylm.model.geometry import geometry_penalty
from babylm.model.layer import ModernBertSmallEncoderLayer
from babylm.model.rotary import RotaryEmbedding


@dataclass
class MaskedLMGeometryOutput(MaskedLMOutput):
    mlm_loss: Optional[torch.Tensor] = None
    geometry_loss: Optional[torch.Tensor] = None
    geometry_centering_loss: Optional[torch.Tensor] = None
    geometry_whitening_loss: Optional[torch.Tensor] = None


class ModernBertSmallPreTrainedModel(PreTrainedModel):
    config_class = ModernBertSmallConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

    def _init_weights(self, module: nn.Module) -> None:
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-2 * std, b=2 * std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, (nn.Embedding, nn.EmbeddingBag)):
            nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-2 * std, b=2 * std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.weight.data.fill_(1.0)
            if module.bias is not None:
                module.bias.data.zero_()


class ModernBertSmallModel(ModernBertSmallPreTrainedModel):
    """The bare encoder: pluggable embeddings -> alternating/uniform-global attention stack ->
    final LayerNorm. RoPE cos/sin and attention masks are computed once per forward pass (one
    global-theta pair, one local-theta pair) and shared across every matching layer."""

    def __init__(self, config: ModernBertSmallConfig) -> None:
        super().__init__(config)
        embedding_cls = get_embedding_class(config.embedding_type)
        self.embeddings = embedding_cls(config)

        attention_types = resolve_attention_types(
            config.num_hidden_layers, config.attention_pattern, config.global_attn_every_n_layers
        )
        self.layers = nn.ModuleList(
            [
                ModernBertSmallEncoderLayer(config, layer_idx, attention_type)
                for layer_idx, attention_type in enumerate(attention_types)
            ]
        )

        head_dim = config.hidden_size // config.num_attention_heads
        self.global_rotary_emb = RotaryEmbedding(head_dim, config.global_rope_theta)
        self.local_rotary_emb = RotaryEmbedding(head_dim, config.local_rope_theta)
        self.local_window_radius = config.local_attention_window // 2
        self.final_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=False)

        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.embeddings

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.embeddings = value

    def forward(
        self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs
    ) -> BaseModelOutput:
        batch, seq_len = input_ids.shape
        device = input_ids.device
        if attention_mask is None:
            attention_mask = torch.ones(batch, seq_len, device=device, dtype=torch.long)

        hidden_states = self.embeddings(input_ids)

        global_mask = build_global_attention_mask(attention_mask)
        local_mask = build_local_attention_mask(attention_mask, self.local_window_radius)
        global_cos, global_sin = self.global_rotary_emb(seq_len, device, hidden_states.dtype)
        local_cos, local_sin = self.local_rotary_emb(seq_len, device, hidden_states.dtype)

        for layer in self.layers:
            if layer.attn.attention_type == "global":
                mask, cos, sin = global_mask, global_cos, global_sin
            else:
                mask, cos, sin = local_mask, local_cos, local_sin
            hidden_states = layer(hidden_states, mask, cos, sin)

        hidden_states = self.final_norm(hidden_states)
        return BaseModelOutput(last_hidden_state=hidden_states)


class ModernBertSmallForMaskedLM(ModernBertSmallPreTrainedModel):
    """MLM head: Dense(d->d, no bias) -> GELU -> LayerNorm(no bias) -> decoder(d->vocab, bias).
    The decoder weight is tied to the embedding module's output weight (when it defines one)."""

    def __init__(self, config: ModernBertSmallConfig) -> None:
        super().__init__(config)
        self.model = ModernBertSmallModel(config)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.head_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=False)
        output_weight = self.model.embeddings.get_output_embedding_weight()
        self.output_proj: Optional[nn.Linear] = None
        self.decoder_bias: Optional[nn.Parameter] = None
        self.structured_output_head: Optional[nn.Module] = None

        if config.tie_word_embeddings:
            self.structured_output_head = self.model.embeddings.build_output_head(config)

        if self.structured_output_head is not None:
            self.decoder = None
            self._tied_weights_keys = []
        elif (
            config.tie_word_embeddings
            and output_weight is not None
            and output_weight.shape[1] != config.hidden_size
        ):
            bottleneck_dim = output_weight.shape[1]
            self.output_proj = nn.Linear(config.hidden_size, bottleneck_dim, bias=False)
            self.decoder = None
            self.decoder_bias = nn.Parameter(torch.zeros(config.vocab_size))
            self._tied_weights_keys = []
        else:
            self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
            self._tied_weights_keys = ["decoder.weight"] if config.tie_word_embeddings else []
        if config.geometry_penalty_mode != "none":
            if config.embedding_type != "compositional" or self.structured_output_head is None:
                raise ValueError(
                    "Geometry penalties require tied compositional embeddings."
                )
            if config.geometry_penalty_weight < 0:
                raise ValueError("geometry_penalty_weight must be non-negative.")
            if config.geometry_sample_size <= 0:
                raise ValueError("geometry_sample_size must be positive.")
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self) -> nn.Module:
        if self.decoder is not None:
            return self.decoder
        if self.output_proj is not None:
            return self.output_proj
        return self.structured_output_head

    def set_output_embeddings(self, value: nn.Module) -> None:
        if self.decoder is not None:
            self.decoder = value
        elif self.output_proj is not None:
            self.output_proj = value
        else:
            self.structured_output_head = value

    def tie_weights(self) -> None:
        if not getattr(self.config, "tie_word_embeddings", True):
            return
        output_weight = self.model.embeddings.get_output_embedding_weight()
        if output_weight is None:
            return
        if self.decoder is None:
            return
        self.decoder.weight = output_weight

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> MaskedLMGeometryOutput:
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = self.head_norm(F.gelu(self.dense(outputs.last_hidden_state)))
        use_geometry = (
            labels is not None
            and self.training
            and self.config.geometry_penalty_mode != "none"
        )
        token_representations = None
        if self.structured_output_head is not None:
            if use_geometry:
                head_output = self.structured_output_head(
                    hidden_states, return_token_representations=True
                )
                logits, token_representations = head_output
            else:
                logits = self.structured_output_head(hidden_states)
        elif self.output_proj is not None:
            output_weight = self.model.embeddings.get_output_embedding_weight()
            logits = F.linear(self.output_proj(hidden_states), output_weight, self.decoder_bias)
        else:
            logits = self.decoder(hidden_states)

        loss = None
        mlm_loss = None
        geometry_loss = None
        geometry_centering_loss = None
        geometry_whitening_loss = None
        if labels is not None:
            mlm_loss = F.cross_entropy(
                logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100
            )
            loss = mlm_loss
            if use_geometry:
                components = geometry_penalty(
                    token_representations,
                    self.config.token_frequencies,
                    self.config.special_token_ids,
                    self.config.geometry_penalty_mode,
                    self.config.geometry_sampling_mode,
                    self.config.geometry_sample_size,
                )
                geometry_loss = components.total
                geometry_centering_loss = components.centering
                geometry_whitening_loss = components.whitening
                loss = mlm_loss + self.config.geometry_penalty_weight * geometry_loss

        return MaskedLMGeometryOutput(
            loss=loss,
            mlm_loss=mlm_loss,
            geometry_loss=geometry_loss,
            geometry_centering_loss=geometry_centering_loss,
            geometry_whitening_loss=geometry_whitening_loss,
            logits=logits,
        )