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, )