Upload folder using huggingface_hub

configuration_modernbert_small.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+from babylm.config.schema import ModelConfig
+
+
+class ModernBertSmallConfig(PretrainedConfig):
+    """HF-compatible config for the from-scratch ModernBERT-Small architecture.
+
+    Bridges to the pydantic `ModelConfig` (used for JSON validation/CLI) via `to_hf_config`,
+    while giving `save_pretrained`/`from_pretrained`/`AutoConfig` compatibility.
+    """
+
+    model_type = "modernbert_small"
+
+    def __init__(
+        self,
+        vocab_size: int = 30522,
+        hidden_size: int = 384,
+        num_hidden_layers: int = 16,
+        num_attention_heads: int = 6,
+        intermediate_size: int = 576,
+        hidden_activation: str = "gelu",
+        attention_pattern: str = "alternating",
+        global_attn_every_n_layers: int = 3,
+        global_rope_theta: float = 160000.0,
+        local_rope_theta: float = 10000.0,
+        local_attention_window: int = 128,
+        norm_eps: float = 1e-5,
+        attention_dropout: float = 0.0,
+        mlp_dropout: float = 0.0,
+        embedding_dropout: float = 0.0,
+        embedding_type: str = "standard",
+        factorized_linear_bottleneck_dim: int = 128,
+        factorized_mlp_bottleneck_dim: int = 128,
+        sttp_rank: int = 8,
+        sttp_num_factors: int = 3,
+        pete_basis_dim: int = 384,
+        compositional_bottleneck_dim: int = 128,
+        compositional_num_buckets: int = 8192,
+        compositional_min_ngram: int = 1,
+        compositional_max_ngram: int = 4,
+        compositional_residual_mode: str = "none",
+        compositional_frequency_tau: float = 100.0,
+        compositional_frequency_shuffle_seed: int = 42,
+        geometry_penalty_mode: str = "none",
+        geometry_penalty_weight: float = 0.01,
+        geometry_sample_size: int = 2048,
+        geometry_sampling_mode: str = "zipfian",
+        compositional_token_strings: Optional[list[str]] = None,
+        token_frequencies: Optional[list[int]] = None,
+        special_token_ids: Optional[list[int]] = None,
+        initializer_range: float = 0.02,
+        pad_token_id: int = 0,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_activation = hidden_activation
+        self.attention_pattern = attention_pattern
+        self.global_attn_every_n_layers = global_attn_every_n_layers
+        self.global_rope_theta = global_rope_theta
+        self.local_rope_theta = local_rope_theta
+        self.local_attention_window = local_attention_window
+        self.norm_eps = norm_eps
+        self.attention_dropout = attention_dropout
+        self.mlp_dropout = mlp_dropout
+        self.embedding_dropout = embedding_dropout
+        self.embedding_type = embedding_type
+        self.factorized_linear_bottleneck_dim = factorized_linear_bottleneck_dim
+        self.factorized_mlp_bottleneck_dim = factorized_mlp_bottleneck_dim
+        self.sttp_rank = sttp_rank
+        self.sttp_num_factors = sttp_num_factors
+        self.pete_basis_dim = pete_basis_dim
+        self.compositional_bottleneck_dim = compositional_bottleneck_dim
+        self.compositional_num_buckets = compositional_num_buckets
+        self.compositional_min_ngram = compositional_min_ngram
+        self.compositional_max_ngram = compositional_max_ngram
+        self.compositional_residual_mode = compositional_residual_mode
+        self.compositional_frequency_tau = compositional_frequency_tau
+        self.compositional_frequency_shuffle_seed = compositional_frequency_shuffle_seed
+        self.geometry_penalty_mode = geometry_penalty_mode
+        self.geometry_penalty_weight = geometry_penalty_weight
+        self.geometry_sample_size = geometry_sample_size
+        self.geometry_sampling_mode = geometry_sampling_mode
+        self.compositional_token_strings = compositional_token_strings
+        self.token_frequencies = token_frequencies
+        self.special_token_ids = special_token_ids or []
+        self.initializer_range = initializer_range
+
+
+def to_hf_config(model_cfg: ModelConfig, pad_token_id: Optional[int] = None) -> ModernBertSmallConfig:
+    """Builds a ModernBertSmallConfig from the pydantic ModelConfig, optionally overriding
+    pad_token_id with the value discovered from a trained tokenizer."""
+    data = model_cfg.model_dump()
+    if pad_token_id is not None:
+        data["pad_token_id"] = pad_token_id
+    return ModernBertSmallConfig(**data)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08073632af2bbb942c5ac108ff4ce6151453f4a948b96811f9a03cb7d064810d
+size 101333640

modeling_modernbert_small.py

@@ -0,0 +1,244 @@
+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,
+        )

optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93e3afcae817960703075fd1b5118b7a6d3a098c6632b287390a0a8f20620108
+size 202520395

rng_state.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06d475bcff9f483e26b631b2347049c0672769aaf37747dbb82794cd30342266
+size 14455

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8b76b363c4f9b730d55cb4d33faed7fa4f8b292e510f6a6a841e72fcadc0e75
+size 1465

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,53 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

trainer_state.json

@@ -0,0 +1 @@
+{"step": 3027, "cumulative_wall_seconds": 220.5930519104004, "cumulative_gpu_seconds": 661.7791557312012, "cumulative_flops": 1.975260676305715e+16, "cumulative_words_seen": 100026227.7631836}