File size: 15,105 Bytes
58050ef 01e34e6 58050ef 01e34e6 58050ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | """
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): # label + trust (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 label+trust (train_loop.Bind, arch "loop2"): verdict-driven trust + experience 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))
|