bind2_0 23.9M build (BabyLM 2026 strict-small training)

config.json

@@ -0,0 +1,29 @@
+{
+  "K": 16,
+  "T": 3,
+  "arch": "bind2_0",
+  "architectures": [
+    "BabyLMForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_babylm.BabyLMConfig",
+    "AutoModel": "modeling_babylm.BabyLMModel",
+    "AutoModelForCausalLM": "modeling_babylm.BabyLMForCausalLM"
+  },
+  "chunk": 32,
+  "core_n": 4,
+  "depth": 12,
+  "dim": 384,
+  "dtype": "float32",
+  "gdn_hd": 72,
+  "gdn_heads": 4,
+  "in_n": 3,
+  "mlp_hidden": 576,
+  "model_type": "babylm",
+  "n_layer": 12,
+  "nhead": 6,
+  "out_n": 3,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.13.0",
+  "vocab_size": 16000
+}

model.safetensors

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

modeling_babylm.py

@@ -0,0 +1,257 @@
+"""
+Self-contained HuggingFace wrapper for the BabyLM entry (LoopLM) and monolith (LM), so the
+models load as a stock AutoModelForCausalLM (trust_remote_code) for babylm-eval / leaderboard.
+Model code is INLINED (no import of train_*.py) so this file is portable on the HF hub.
+The ACTIVE class defs (LoopLMv2/Bind2 for arch "loop2", LM for the monolith) are byte-for-byte
+the current training defs (train_loop.py / train_stage1.py) so state_dicts load exactly; the
+legacy v1 defs (LoopLM/Bind) are retained ONLY to load the already-published v1 bypass
+checkpoint (paper §4b diagnostic) and no longer exist in train_loop.py. forward() runs the whole loop inside a standard causal pass and
+returns CausalLMOutput(logits, loss); empty-context, stateless across examples.
+
+BabyLMModel (AutoModel entry) exists for the GLUE finetuning pipeline, which pools
+last_hidden_state through its own classifier head. attention_mask is honored only on that
+path (padded batches); the causal-LM path is unchanged — attn_mask=None reproduces the
+exact zero-shot behavior the published eval numbers came from.
+"""
+import math, torch, torch.nn as nn, torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutput
+
+def build_rope(T, D, device, base=10000.0):
+    inv = 1.0/(base**(torch.arange(0,D,2,device=device).float()/D)); t = torch.arange(T,device=device).float()
+    f = torch.outer(t, inv); emb = torch.cat([f, f], dim=-1); return emb.cos(), emb.sin()
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1); return torch.cat((-x2, x1), dim=-1)
+def apply_rope(x, cos, sin):
+    return x*cos[None,None] + rotate_half(x)*sin[None,None]
+
+class Attn(nn.Module):
+    def __init__(self, d, nh):
+        super().__init__(); self.nh=nh; self.hd=d//nh
+        self.qkv=nn.Linear(d,3*d,bias=False); self.o=nn.Linear(d,d,bias=False)
+    def forward(self, x, cos, sin, attn_mask=None):
+        B,T,D=x.shape; qkv=self.qkv(x).view(B,T,3,self.nh,self.hd).permute(2,0,3,1,4)
+        q,k,v=qkv[0],qkv[1],qkv[2]; q=apply_rope(q,cos,sin); k=apply_rope(k,cos,sin)
+        if attn_mask is None: o=F.scaled_dot_product_attention(q,k,v,is_causal=True)
+        else: o=F.scaled_dot_product_attention(q,k,v,attn_mask=attn_mask)
+        return self.o(o.transpose(1,2).reshape(B,T,D))
+class SwiGLU(nn.Module):
+    def __init__(self, d, h):
+        super().__init__(); self.w1=nn.Linear(d,h,bias=False); self.w3=nn.Linear(d,h,bias=False); self.w2=nn.Linear(h,d,bias=False)
+    def forward(self, x): return self.w2(F.silu(self.w1(x))*self.w3(x))
+class Block(nn.Module):
+    def __init__(self, d, nh, h):
+        super().__init__(); self.n1=nn.RMSNorm(d); self.attn=Attn(d,nh); self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d,h)
+    def forward(self, x, cos, sin, attn_mask=None):
+        x=x+self.attn(self.n1(x),cos,sin,attn_mask); return x+self.mlp(self.n2(x))
+
+class LM(nn.Module):                       # monolith (train_stage1.LM)
+    def __init__(self, vocab, d=384, nl=12, nh=6):
+        super().__init__(); h=((int(8/3*d)+63)//64)*64
+        self.emb=nn.Embedding(vocab,d); self.blocks=nn.ModuleList([Block(d,nh,h) for _ in range(nl)])
+        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
+        self.d=d; self.nh=nh
+    def hidden(self, ids, attn_mask=None):
+        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
+        for b in self.blocks: h=b(h,cos,sin,attn_mask)
+        return self.nf(h)
+    def forward(self, ids): return self.head(self.hidden(ids))
+
+class Bind(nn.Module):                     # 想 + 行 (train_loop.Bind)
+    def __init__(self, d, K=16, dr=64):
+        super().__init__(); self.role=nn.Linear(d,K,bias=False); self.R=nn.Parameter(torch.randn(K,dr)*0.02)
+        self.up=nn.Linear(dr,d,bias=False); self.trust=nn.Linear(d,1)
+    def forward(self, h):
+        a=torch.softmax(self.role(h),dim=-1); lab=a@self.R; tau=torch.sigmoid(self.trust(h)); return h+tau*self.up(lab)
+class LoopLM(nn.Module):                    # entry (train_loop.LoopLM)
+    def __init__(self, vocab, d=384, in_n=3, core_n=4, out_n=3, nh=6, T=3, K=16):
+        super().__init__(); hdim=((int(8/3*d)+63)//64)*64
+        self.emb=nn.Embedding(vocab,d)
+        self.inb=nn.ModuleList([Block(d,nh,hdim) for _ in range(in_n)])
+        self.core=nn.ModuleList([Block(d,nh,hdim) for _ in range(core_n)])
+        self.outb=nn.ModuleList([Block(d,nh,hdim) for _ in range(out_n)])
+        self.bind=Bind(d,K); self.vhead=nn.Linear(d,1)
+        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
+        self.d=d; self.nh=nh; self.T=T
+    def hidden(self, ids, attn_mask=None):
+        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
+        for b in self.inb: h=b(h,cos,sin,attn_mask)
+        for _ in range(self.T):
+            z=self.bind(h); h2=z
+            for b in self.core: h2=b(h2,cos,sin,attn_mask)
+            v=torch.sigmoid(self.vhead(h2)); h=h+(1.0-v)*(h2-h)
+        for b in self.outb: h=b(h,cos,sin,attn_mask)
+        return self.nf(h)
+    def forward(self, ids): return self.head(self.hidden(ids))
+
+class Bind2(nn.Module):                    # v2 想+行 (train_loop.Bind, arch "loop2"): 受-driven trust + 熏習 prior + role-slice re-stamp
+    def __init__(self, d, K=16, dr=64):
+        super().__init__()
+        self.dr = dr
+        self.role = nn.Linear(d, K, bias=False)
+        self.role_scale = nn.Parameter(torch.ones(1))
+        self.R = nn.Parameter(torch.randn(K, dr) * 0.02)
+        self.trust = nn.Linear(d, 1)
+        self.v_gain = nn.Parameter(torch.zeros(1))
+        self.vasana = nn.Parameter(torch.zeros(K))
+    def forward(self, h, v_prev):
+        a = torch.softmax(self.role_scale * self.role(h), dim=-1)
+        lab = a @ self.R
+        tau = torch.sigmoid(self.trust(h) + (a @ self.vasana)[..., None] + self.v_gain * (0.5 - v_prev))
+        s = h[..., -self.dr:]
+        return torch.cat([h[..., :-self.dr], (1.0 - tau) * s + tau * lab], dim=-1), a, tau
+
+class LoopLMv2(nn.Module):                  # entry v2 (train_loop.LoopLM, arch "loop2")
+    def __init__(self, vocab, d=384, in_n=3, core_n=4, out_n=3, nh=6, T=3, K=16):
+        super().__init__(); hdim=((int(8/3*d)+63)//64)*64
+        self.emb=nn.Embedding(vocab,d)
+        self.inb=nn.ModuleList([Block(d,nh,hdim) for _ in range(in_n)])
+        self.core=nn.ModuleList([Block(d,nh,hdim) for _ in range(core_n)])
+        self.outb=nn.ModuleList([Block(d,nh,hdim) for _ in range(out_n)])
+        self.bind=Bind2(d,K); self.vhead=nn.Linear(d,1)
+        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
+        self.d=d; self.nh=nh; self.T=T
+    def hidden(self, ids, attn_mask=None):
+        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
+        for b in self.inb: h=b(h,cos,sin,attn_mask)
+        v=torch.full_like(h[..., :1], 0.5)
+        for _ in range(self.T):
+            z,a,tau=self.bind(h,v); h2=z
+            for b in self.core: h2=b(h2,cos,sin,attn_mask)
+            v=torch.sigmoid(self.vhead(h2)); h=h2          # state flows through the loop (no bypass)
+        for b in self.outb: h=b(h,cos,sin,attn_mask)
+        return self.nf(h)
+    def forward(self, ids): return self.head(self.hidden(ids))
+
+# --- delta-rule + forced-bottleneck (arch "bind2_0"); class defs byte-for-byte from modeling_bind2_0.py
+#     (train_bind2_0_babylm.py) so state_dicts load exactly. fla is imported lazily inside GDNBlock so
+#     this module still imports without fla for the mono/loop2 paths. ---
+class ChunkedAttn(nn.Module):
+    """Forced bottleneck: causal attention restricted to within non-overlapping chunks of size C."""
+    def __init__(self, d, nh, chunk):
+        super().__init__()
+        self.nh=nh; self.hd=d//nh; self.chunk=chunk
+        self.qkv=nn.Linear(d,3*d,bias=False); self.o=nn.Linear(d,d,bias=False)
+    def forward(self, x, cos, sin):
+        B,T,D=x.shape
+        qkv=self.qkv(x).view(B,T,3,self.nh,self.hd).permute(2,0,3,1,4)
+        q,k,v=qkv[0],qkv[1],qkv[2]
+        q=apply_rope(q,cos,sin); k=apply_rope(k,cos,sin)
+        idx=torch.arange(T,device=x.device)
+        same=(idx[:,None]//self.chunk)==(idx[None,:]//self.chunk)
+        causal=idx[:,None]>=idx[None,:]
+        keep=same&causal
+        mask=torch.zeros(T,T,device=x.device,dtype=q.dtype).masked_fill(~keep,float("-inf"))
+        o=F.scaled_dot_product_attention(q,k,v,attn_mask=mask)
+        return self.o(o.transpose(1,2).reshape(B,T,D))
+class GDNBlock(nn.Module):
+    def __init__(self, d, idx, mlp_hidden, gdn_heads=4, gdn_hd=72):
+        super().__init__()
+        from fla.layers import GatedDeltaNet    # lazy: only bind2_0 needs fla
+        self.n1=nn.RMSNorm(d)
+        self.gdn=GatedDeltaNet(hidden_size=d, num_heads=gdn_heads, head_dim=gdn_hd, layer_idx=idx)
+        self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d, mlp_hidden)
+    def forward(self, x):
+        m=self.gdn(self.n1(x))[0]   # fla returns (output, attn, cache)
+        x=x+m
+        return x+self.mlp(self.n2(x))
+class AttnBlock(nn.Module):
+    def __init__(self, d, nh, chunk, mlp_hidden):
+        super().__init__()
+        self.n1=nn.RMSNorm(d); self.attn=ChunkedAttn(d,nh,chunk)
+        self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d,mlp_hidden)
+    def forward(self, x, cos, sin):
+        x=x+self.attn(self.n1(x),cos,sin)
+        return x+self.mlp(self.n2(x))
+class Bind2_0LM(nn.Module):                 # delta-rule + forced-bottleneck (modeling_bind2_0.Bind2_0LM)
+    def __init__(self, vocab, d=384, depth=12, nh=6, chunk=32, mlp_hidden=576, gdn_heads=4, gdn_hd=72):
+        super().__init__()
+        self.emb=nn.Embedding(vocab,d)
+        self.kinds=["attn" if (i+1)%4==0 else "gdn" for i in range(depth)]   # 3:1 GDN:attn
+        self.blocks=nn.ModuleList([
+            GDNBlock(d,i,mlp_hidden,gdn_heads,gdn_hd) if k=="gdn" else AttnBlock(d,nh,chunk,mlp_hidden)
+            for i,k in enumerate(self.kinds)])
+        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
+        self.d=d; self.nh=nh; self.chunk=chunk
+    def hidden(self, ids, attn_mask=None):   # attn_mask unused: chunked attn carries its own intra-chunk
+        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)   # mask (pad-mask
+        for blk,k in zip(self.blocks,self.kinds):                                          # for GLUE is TODO,
+            h=blk(h) if k=="gdn" else blk(h,cos,sin)                                        # zero-shot unaffected)
+        return self.nf(h)
+    def forward(self, ids): return self.head(self.hidden(ids))
+
+def _build_backbone(config):
+    if config.arch == "bind2_0":
+        return Bind2_0LM(config.vocab_size, config.dim, config.depth, config.nhead,
+                         chunk=config.chunk, mlp_hidden=config.mlp_hidden,
+                         gdn_heads=config.gdn_heads, gdn_hd=config.gdn_hd)
+    if config.arch == "loop2":
+        return LoopLMv2(config.vocab_size, config.dim, config.in_n, config.core_n,
+                        config.out_n, config.nhead, config.T, config.K)
+    if config.arch == "loop":
+        return LoopLM(config.vocab_size, config.dim, config.in_n, config.core_n,
+                      config.out_n, config.nhead, config.T, config.K)
+    return LM(config.vocab_size, config.dim, config.n_layer, config.nhead)
+
+class BabyLMConfig(PretrainedConfig):
+    model_type = "babylm"
+    # the GLUE finetuning classifier reads config.hidden_size
+    attribute_map = {"hidden_size": "dim", "num_attention_heads": "nhead", "num_hidden_layers": "n_layer"}
+    def __init__(self, arch="loop", vocab_size=16000, dim=384, in_n=3, core_n=4, out_n=3,
+                 T=3, K=16, nhead=6, n_layer=12,
+                 depth=12, chunk=32, mlp_hidden=576, gdn_heads=4, gdn_hd=72, **kw):
+        self.arch=arch; self.vocab_size=vocab_size; self.dim=dim; self.in_n=in_n; self.core_n=core_n
+        self.out_n=out_n; self.T=T; self.K=K; self.nhead=nhead; self.n_layer=n_layer
+        self.depth=depth; self.chunk=chunk; self.mlp_hidden=mlp_hidden; self.gdn_heads=gdn_heads; self.gdn_hd=gdn_hd
+        super().__init__(**kw)
+
+class BabyLMForCausalLM(PreTrainedModel):
+    config_class = BabyLMConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = _build_backbone(config)
+        # Untie the LM head for a clean HF save (no shared tensors). Inference-equivalent: the head
+        # weight is loaded from the checkpoint, which equals the tied embedding used at train time.
+        self.backbone.head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        self.config.tie_word_embeddings = False
+        self.post_init()
+    def tie_weights(self, *args, **kwargs):
+        pass  # head intentionally untied for export
+    def get_input_embeddings(self): return self.backbone.emb
+    def set_input_embeddings(self, v): self.backbone.emb = v
+    def get_output_embeddings(self): return self.backbone.head
+    def forward(self, input_ids=None, labels=None, attention_mask=None, **kw):
+        logits = self.backbone(input_ids)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)).float(), labels[:, 1:].reshape(-1))
+        return CausalLMOutput(loss=loss, logits=logits)
+
+def padding_causal_mask(attention_mask):
+    # bool SDPA mask (B,1,T,T): attend where causal AND the key is a real (non-pad) token.
+    # Pad-query rows would be fully masked (softmax NaN) with left padding, so the diagonal
+    # stays open; their outputs are finite and get zero weight from every real query.
+    B, T = attention_mask.shape; dev = attention_mask.device
+    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=dev))
+    m = causal[None, None] & attention_mask.to(torch.bool)[:, None, None, :]
+    return m | torch.eye(T, dtype=torch.bool, device=dev)[None, None]
+
+class BabyLMModel(PreTrainedModel):
+    """AutoModel entry (base model, no LM head applied) for the GLUE finetuning pipeline.
+    Same backbone module tree as BabyLMForCausalLM so the exported checkpoint loads key-for-key."""
+    config_class = BabyLMConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = _build_backbone(config)
+        self.backbone.head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        self.config.tie_word_embeddings = False
+        self.post_init()
+    def tie_weights(self, *args, **kwargs):
+        pass  # head intentionally untied for export
+    def get_input_embeddings(self): return self.backbone.emb
+    def set_input_embeddings(self, v): self.backbone.emb = v
+    def forward(self, input_ids=None, attention_mask=None, **kw):
+        attn_mask = None
+        if attention_mask is not None and not bool(attention_mask.all()):
+            attn_mask = padding_causal_mask(attention_mask)
+        return BaseModelOutput(last_hidden_state=self.backbone.hidden(input_ids, attn_mask))

requirements_pins.txt

@@ -0,0 +1,25 @@
+# Exact dependency versions the bind2_0 checkpoints were trained/exported with.
+# Source: local conda env `babylm-smoke` (C:/Users/yulin/.conda/envs/babylm-smoke),
+# queried via pip on 2026-07-15.
+#
+# Python: 3.11.15
+#
+# NOTE on fla (flash-linear-attention): the inlined HF modeling code
+# (modeling_babylm.py) DOES import it — lazily, inside GDNBlock:
+#   `from fla.layers import GatedDeltaNet`  (only executed when arch == "bind2_0")
+# Since the bind2_0 exports instantiate GDN blocks, fla IS required at runtime
+# to load/run these checkpoints. Installed from PyPI as release 0.5.1
+# (no commit-pin / direct-URL metadata present in the env; pulls fla-core 0.5.1).
+#
+# NOTE on triton: the env uses the Windows fork `triton-windows`
+# (github.com/woct0rdho/triton-windows); on Linux use the matching upstream
+# `triton` that your torch build requires.
+# torch build is CUDA 12.6 (`+cu126`); pick the equivalent build for your platform.
+
+torch==2.12.1+cu126
+transformers==5.13.0
+triton-windows==3.7.1.post27
+flash-linear-attention==0.5.1
+fla-core==0.5.1
+safetensors==0.8.0
+numpy==2.4.6

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}