Upload RecGPT-10M main checkpoint

config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "RecGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_recgpt.RecGPTConfig",
+    "AutoModelForCausalLM": "modeling_recgpt.RecGPTForCausalLM"
+  },
+  "dtype": "float32",
+  "embedding_size": 192,
+  "head_dim": 64,
+  "hidden_size": 768,
+  "intermediate_size": 12288,
+  "is_decoder": true,
+  "max_position_embeddings": 1024,
+  "model_type": "recgpt",
+  "num_heads": 12,
+  "pad_token_id": 0,
+  "recursive_depth": 16,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.9.0",
+  "use_cache": false,
+  "vocab_size": 32768
+}

model.safetensors

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

modeling_recgpt.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers import initialization as init
+from transformers.modeling_outputs import CausalLMOutput
+
+
+class RecGPTConfig(PretrainedConfig):
+    model_type = "recgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 32768,
+        hidden_size: int = 640,
+        embedding_size: int = 192, # Allows for factorized embeddings, only makes sense at babylm scale.
+        head_dim: int = 64,
+        intermediate_size: int = 10240,
+        recursive_depth: int = 16,
+        max_position_embeddings: int = 1024,
+        pad_token_id: int = 0, # Padding is determined by segment_ids, this is only used for embeddings/HF metadata.
+        tie_word_embeddings: bool = False, # Tied embeddings greatly hurt performance for recursive models.
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        if hidden_size % head_dim != 0:
+            raise ValueError("hidden_size must be divisible by head_dim.")
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.embedding_size = embedding_size
+        self.head_dim = head_dim
+        self.num_heads = hidden_size // head_dim
+        self.intermediate_size = intermediate_size
+        self.recursive_depth = recursive_depth
+        self.max_position_embeddings = max_position_embeddings
+        self.is_decoder = True
+        self.use_cache = False
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6, use_bias: bool = False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size)) if use_bias else None
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.rms_norm(x, (x.size(-1),), self.weight, self.eps)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, max_position_embeddings: int, theta: float):
+        super().__init__()
+        if head_dim % 2 != 0:
+            raise ValueError("RoPE requires an even head dimension.")
+
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.theta = theta
+
+        # We register these as buffers to ensure they get moved to device together with the model.
+        self.register_buffer("cos", torch.empty(max_position_embeddings, head_dim // 2), persistent=False)
+        self.register_buffer("sin", torch.empty(max_position_embeddings, head_dim // 2), persistent=False)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, self.head_dim, 2, device=self.cos.device, dtype=torch.float32)
+                / self.head_dim
+            )
+        )
+        positions = torch.arange(self.max_position_embeddings, device=self.cos.device, dtype=torch.float32)
+        freqs = torch.outer(positions, inv_freq)
+        init.copy_(self.cos, freqs.cos())
+        init.copy_(self.sin, freqs.sin())
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+        cos = self.cos[position_ids].unsqueeze(2).to(dtype=x.dtype)
+        sin = self.sin[position_ids].unsqueeze(2).to(dtype=x.dtype)
+        x_even = x[..., 0::2]
+        x_odd = x[..., 1::2]
+
+        out = torch.empty_like(x)
+        out[..., 0::2] = x_even * cos - x_odd * sin
+        out[..., 1::2] = x_odd * cos + x_even * sin
+        return out
+
+
+def _select_flex_backend(device: torch.device) -> str:
+    if device.type != "cuda":
+        raise RuntimeError("RecGPT attention requires a CUDA/ROCm accelerator.")
+    if torch.version.hip is None:
+        major, _ = torch.cuda.get_device_capability(device)
+        if major >= 9:
+            return "FLASH"
+    return "TRITON"
+
+ROPE_THETA = 10000.0
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: RecGPTConfig):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_heads
+        self.head_dim = config.head_dim
+
+        # One fused projection for QKV.
+        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
+        self.out = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        # Zero init (Idea from modded-nanogpt speedrun, empirically seems to work well).
+        nn.init.zeros_(self.out.weight)
+        self.rope = RotaryEmbedding(self.head_dim, config.max_position_embeddings, ROPE_THETA)
+
+        # Gated Attention (https://arxiv.org/pdf/2505.06708)
+        # SDPAHeadwiseGate: per-head sigmoid gate applied to attention output.
+        # Empirically, attention gating should benefit us since we don't use any <|BOS|> token during training,
+        # meaning the model has no attention sinks. GA should reduce the need for attention sinks.
+        self.gate = nn.Linear(config.hidden_size, self.num_heads, bias=True)
+        nn.init.zeros_(self.gate.weight)
+
+        # Since the gated attention paper finds that the model converges toward a more sparse gate,
+        # we initialize the gate bias with 0.0 (so the sigmoid of the bias is 0.5).
+        # 0.5 is right in the middle, not too high to start with default behavior, not too low to enforce sparsity early on.
+        # TODO: Rewrite this comment more clearly
+        # TODO: Re-consider if bias is even necessary
+        nn.init.constant_(self.gate.bias, 0.0)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        block_mask,
+        backend: str,
+    ) -> torch.Tensor:
+        from torch.nn.attention.flex_attention import flex_attention
+
+        batch_size, seq_len, hidden_size = x.shape
+        # Project to QKV and reshape to [B, T, 3, H, D].
+        qkv = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+
+        # Unparameterized QK norm. We used parameterized QK norms in the old repo,
+        # but this lean version keeps them fixed for now.
+        q = F.rms_norm(q, (self.head_dim,))
+        k = F.rms_norm(k, (self.head_dim,))
+
+        # Pick cos/sin for each token position, then broadcast over heads.
+        q = self.rope(q, position_ids).transpose(1, 2)
+        k = self.rope(k, position_ids).transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        y = flex_attention(q, k, v, block_mask=block_mask, kernel_options={"BACKEND": backend})
+        gate = torch.sigmoid(self.gate(x)).view(batch_size, seq_len, self.num_heads, 1)
+        y = y.transpose(1, 2) * gate
+        y = y.contiguous().view(batch_size, seq_len, hidden_size)
+        return self.out(y)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: RecGPTConfig):
+        super().__init__()
+        self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        # Standard dense ReLU^2 MLP with zero init output
+        # (Idea from modded-nanogpt speedrun, empirically seems to work well).
+        nn.init.zeros_(self.down.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.relu(self.up(x)).square())
+
+
+class RecGPTForCausalLM(PreTrainedModel):
+    config_class = RecGPTConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: RecGPTConfig):
+        super().__init__(config)
+        self.use_factorized = config.embedding_size != config.hidden_size
+
+        # Factorized Embeddings (https://arxiv.org/pdf/1909.11942)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)
+        if self.use_factorized:
+            self.e_to_h = nn.Linear(config.embedding_size, config.hidden_size, bias=False)
+            self.h_to_e = nn.Linear(config.hidden_size, config.embedding_size, bias=False)
+        if not config.tie_word_embeddings:
+            self.lm_head = nn.Linear(config.embedding_size, config.vocab_size, bias=False)
+
+        # This is the main recursive model idea: attention and MLP weights are
+        # reused at every depth. The norms are depth-specific.
+        self.attn = SelfAttention(config)
+        self.mlp = MLP(config)
+        self.attn_norms = nn.ModuleList([RMSNorm(config.hidden_size, use_bias=True) for _ in range(config.recursive_depth)])
+        self.mlp_norms = nn.ModuleList([RMSNorm(config.hidden_size, use_bias=True) for _ in range(config.recursive_depth)])
+        self.final_norm = RMSNorm(config.hidden_size)
+
+        self.post_init()
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, RotaryEmbedding):
+            module.reset_parameters()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return getattr(self, "lm_head", None)
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward( # Everything expected in shape [batch, seq_len]
+        self,
+        input_ids: torch.Tensor,
+        segment_ids: Optional[torch.Tensor] = None, # We use segment_ids for our packed training
+        attention_mask: Optional[torch.Tensor] = None, # But we also support attention_mask for compatibility with HF transformers stack.
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_hidden_and_embed: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutput | tuple[torch.Tensor, ...]:
+        # Training uses one packed, block-masked sequence per microbatch, so B is always 1.
+        # We keep the batch dimension because HF expects it.
+        if input_ids.dim() != 2:
+            raise ValueError("input_ids must have shape [batch, seq_len].")
+        if segment_ids is not None and segment_ids.shape != input_ids.shape:
+            raise ValueError("segment_ids must match input_ids shape.")
+        if attention_mask is not None and segment_ids is None: # Compatibility with HF transformers stack. Ignored if segment_ids are provided.
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError("attention_mask must match input_ids shape when segment_ids is not provided.")
+            # Assumes standard binary attention mask where 1 indicates a real token and 0 indicates padding. 
+            # Converts to segment_ids where padding is -1 and real tokens are >= 0.
+            assert ((attention_mask == 0) | (attention_mask == 1)).all()
+            segment_ids = attention_mask.to(device=input_ids.device, dtype=torch.long) - 1
+        if segment_ids is None:
+            segment_ids = torch.zeros_like(input_ids)
+
+        assert segment_ids.device == input_ids.device # If segment_ids are passed explicitly, they should be on the same device.
+
+        # Construct position_ids from segment_ids
+        valid = segment_ids >= 0
+        seq_positions = torch.arange(input_ids.size(1), device=input_ids.device, dtype=torch.long).unsqueeze(0)
+
+        # This marks where a valid segment starts. A token is a segment start if:
+        # - It's valid
+        # - Either it is the first token, or its segment_id differs from the previous token.
+        segment_starts = valid & torch.cat(
+            [
+                torch.ones(segment_ids.size(0), 1, device=input_ids.device, dtype=torch.bool), # First token
+                segment_ids[:, 1:] != segment_ids[:, :-1], # Segment ID changes
+            ],
+            dim=1,
+        )
+
+        # Put each segment start's absolute index at the start token, then cummax fills the latest start index across that segment.
+        # Subtracting it from the absolute token index gives local positions per segment.
+        segment_start_positions = torch.where(segment_starts, seq_positions, 0).cummax(dim=-1).values
+        position_ids = (seq_positions - segment_start_positions).masked_fill(~valid, 0)
+
+        from torch.nn.attention.flex_attention import create_block_mask
+
+        batch_size, seq_len = input_ids.shape
+
+        def mask_mod(b, h, q_idx, kv_idx):
+            """
+            FlexAttention calls mask_mod with scalar/block index tensors and uses the
+            result to build a block-sparse attention mask. This is the only place where
+            document boundaries are enforced.
+            
+            segment_ids[b, t] >= 0 means a real token. segment_ids[b, t] == -1 means
+            padding. Tokens can only attend causally within the same segment, so packed
+            documents in the same row still have hard attention boundaries.
+            """
+            
+            valid = segment_ids[b, q_idx] >= 0
+            same_segment = segment_ids[b, q_idx] == segment_ids[b, kv_idx]
+            causal = kv_idx <= q_idx
+            return valid & same_segment & causal
+
+        # BlockMask is built once per forward and reused at every recursive depth.
+        block_mask = create_block_mask(
+            mask_mod,
+            B=batch_size,
+            H=self.config.num_heads,
+            Q_LEN=seq_len,
+            KV_LEN=seq_len,
+            device=input_ids.device,
+        )
+        backend = _select_flex_backend(input_ids.device)
+
+        x = self.embed_tokens(input_ids)
+        if self.use_factorized:
+            x = self.e_to_h(x)
+        if return_hidden_and_embed:
+            e = x
+        for attn_norm, mlp_norm in zip(self.attn_norms, self.mlp_norms):
+            # We do pre-norm and QK norm.
+            # We used to do a Gemma 3 style post-norm, but removed it to improve stability
+            # and keep the residual stream norm in check. Seems to work fine.
+            # Update: Tried KEEL norm paper with residual scaling, it hurt performance.
+            x = x + self.attn(attn_norm(x), position_ids, block_mask, backend)
+            x = x + self.mlp(mlp_norm(x))
+
+        x = self.final_norm(x) # Final normalized hidden state before output projection.
+        if self.use_factorized:
+            y = self.h_to_e(x)
+        else:
+            y = x
+        if hasattr(self, "lm_head"):
+            logits = self.lm_head(y)
+        else:
+            logits = F.linear(y, self.embed_tokens.weight)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100)
+
+        use_return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
+        if use_return_dict:
+            return CausalLMOutput(loss=loss, logits=logits, hidden_states=(x,) if output_hidden_states else None)
+        if loss is None:
+            return (logits,) # [batch, seq_len, vocab_size]
+        if return_hidden_and_embed:
+            return (loss, logits, x, e) # logits: [batch, seq_len, vocab_size], x: [batch, seq_len, hidden_size], e: [batch, seq_len, hidden_size]
+        return (loss, logits)

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "backend": "tokenizers",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend"
+}

train_config.json

@@ -0,0 +1,74 @@
+{
+  "dataset": "s33c67-10m.parquet",
+  "tokenizer": "s33c67-10m-bpe",
+  "run_name": "recgpt-10m-submission",
+  "seed": 0,
+  "data_seed": 0,
+  "microbatch_tok": 32768,
+  "total_batch_tok": 32768,
+  "sequence_len": 256,
+  "epochs": 10,
+  "checkpoint_track": "strict-small",
+  "max_tokens": -1,
+  "lr_embed": 0.005,
+  "lr_block": 0.02,
+  "min_lr": 0.0,
+  "wd_adam": 0.005,
+  "wd_muon": 0.1,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.997,
+  "muon_momentum": 0.95,
+  "warmup_ratio": 0.0,
+  "cooldown_ratio": 0.2,
+  "max_grad_norm": 2.0,
+  "nl_mult": 0.01,
+  "nl_depth": 2,
+  "nl_hidden": -1,
+  "nl_intermediate": 5120,
+  "nl_lr": 0.004,
+  "nl_wd": 0.01,
+  "nl_momentum": 0.95,
+  "torch_compile": true,
+  "use_wandb": true,
+  "wandb_project": "bblm26-recgpt",
+  "log_every": 10,
+  "model_config": {
+    "transformers_version": "5.9.0",
+    "architectures": [
+      "RecGPTForCausalLM"
+    ],
+    "output_hidden_states": false,
+    "return_dict": true,
+    "dtype": "float32",
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "problem_type": null,
+    "_name_or_path": "",
+    "pad_token_id": 0,
+    "tie_word_embeddings": false,
+    "vocab_size": 32768,
+    "hidden_size": 768,
+    "embedding_size": 192,
+    "head_dim": 64,
+    "num_heads": 12,
+    "intermediate_size": 12288,
+    "recursive_depth": 16,
+    "max_position_embeddings": 1024,
+    "is_decoder": true,
+    "use_cache": false,
+    "auto_map": {
+      "AutoConfig": "modeling_recgpt.RecGPTConfig",
+      "AutoModelForCausalLM": "modeling_recgpt.RecGPTForCausalLM"
+    },
+    "model_type": "recgpt",
+    "output_attentions": false
+  }
+}