Upload folder using huggingface_hub

config.json

@@ -0,0 +1,33 @@
+{
+  "architectures": [
+    "ModernBertSmallForMaskedLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_pattern": "alternating",
+  "auto_map": {
+    "AutoConfig": "configuration_modernbert_small.ModernBertSmallConfig",
+    "AutoModelForMaskedLM": "modeling_modernbert_small.ModernBertSmallForMaskedLM"
+  },
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "embedding_type": "factorized_linear",
+  "factorized_linear_bottleneck_dim": 128,
+  "factorized_mlp_bottleneck_dim": 128,
+  "global_attn_every_n_layers": 3,
+  "global_rope_theta": 160000.0,
+  "hidden_activation": "gelu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 576,
+  "local_attention_window": 128,
+  "local_rope_theta": 10000.0,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert_small",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 6,
+  "num_hidden_layers": 16,
+  "pad_token_id": 0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "vocab_size": 30522
+}

configuration_modernbert_small.py

@@ -0,0 +1,72 @@
+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,
+        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.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:39335d6881ad2f82d46f33216ff40c73c2141c6c1069e1eae8904533cb6e0939
+size 143701720

modeling_modernbert_small.py

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

optimizer.pt

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

rng_state.pt

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

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c59d98e7f49d748953e6dfad099811cb493aa2b1d98d54450091e6a3176535ca
+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": 9069, "cumulative_wall_seconds": 18772.55534029007, "cumulative_gpu_seconds": 18772.55534029007, "cumulative_flops": 1.972650487106765e+16, "cumulative_words_seen": 100001352.76000977}