from __future__ import annotations 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.layer import ModernBertSmallEncoderLayer from babylm.model.rotary import RotaryEmbedding 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.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 if ( 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 [] 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: return self.decoder if self.decoder is not None else self.output_proj def set_output_embeddings(self, value: nn.Module) -> None: if self.decoder is not None: self.decoder = value else: self.output_proj = 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, ) -> MaskedLMOutput: outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) hidden_states = self.head_norm(F.gelu(self.dense(outputs.last_hidden_state))) if 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 if labels is not None: loss = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100 ) return MaskedLMOutput(loss=loss, logits=logits)