File size: 38,641 Bytes
f7bd29e | 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 | {
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Shared, executable implementation for Part 1\nAll five model notebooks execute these definition cells. Data preparation, experiments and outputs remain in notebooks. No Python module is required. The assignment prohibits LLM assistance with the submitted writeup; markdown here documents software and measurement protocols, not a submission report.\n\nSources: [assignment](../../ANLP-Assignment-2.pdf), [dataset](https://huggingface.co/datasets/belumind/en-vi-ja-curated-500k-triplets), [Switch Transformer](https://arxiv.org/abs/2101.03961), [PyTorch SDPA](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html), [SacreBLEU](https://github.com/mjpost/sacrebleu).\n",
"id": "73747e3498e8"
},
{
"cell_type": "code",
"metadata": {
"tags": [
"definitions"
]
},
"source": "import os, json, math, time, random, hashlib, contextlib, platform, unicodedata, gc, io\nfrom pathlib import Path\nfrom dataclasses import dataclass, asdict, replace\nfrom collections import Counter\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom tokenizers import Tokenizer, models, normalizers, pre_tokenizers, decoders, trainers\nimport sacrebleu\nfrom IPython.display import display, Markdown\n\ndef find_root():\n here=Path.cwd().resolve()\n for p in [here,*here.parents]:\n if (p/'shared'/'implementation.ipynb').exists(): return p\n if (here/'part 1'/'shared'/'implementation.ipynb').exists(): return here/'part 1'\n raise FileNotFoundError('Run with the working directory inside part 1.')\nROOT=find_root()\nos.environ.setdefault('HF_HOME', str(ROOT/'.cache'/'huggingface'))\nos.environ.setdefault('TOKENIZERS_PARALLELISM','false')\nDEVICE=torch.device('cuda' if torch.cuda.is_available() else 'cpu')\ntorch.set_num_threads(min(4,os.cpu_count() or 1))\nPAD,BOS,EOS,SEP=0,1,2,3\nLANGS=('vi','ja')\nDATASET='belumind/en-vi-ja-curated-500k-triplets'\n\ndef seed_all(seed=42):\n random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)\n if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)\n\ndef sha(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest()\ndef atomic_torch(obj,path):\n path=Path(path); temp=path.with_suffix(path.suffix+'.tmp'); torch.save(obj,temp); temp.replace(path)\ndef amp():\n return torch.autocast('cuda',dtype=torch.float16) if DEVICE.type=='cuda' else contextlib.nullcontext()\n\n@dataclass\nclass Config:\n variant:int=1\n vocab:int=16000\n d:int=512\n heads:int=8\n layers:int=6\n hidden:int=2048\n context:int=256\n dropout:float=0.1\n seed:int=42\n token_budget:int=60_000_000\n batch_tokens:int=8192\n batch_examples:int=128\n lr:float=0.0006\n weight_decay:float=0.1\n balance_coef:float=0.01\n z_coef:float=0.001\n eval_interval:int=5_000_000\n val_examples:int=2048\n max_new_tokens:int=160\n profile:str='full'\n\nNAMES={1:'Dense two-layer MLP',2:'4 experts; top 1; equal total parameters',\n 3:'4 experts; top 2; equal total parameters',\n 4:'1 always-active shared + 3 routed; top 1; equal total parameters',\n 5:'4 experts; top 2; equal active parameters'}\n",
"execution_count": null,
"outputs": [],
"id": "8527fab62a04"
},
{
"cell_type": "markdown",
"metadata": {},
"source": "## Parameter accounting and routing\nThe dense FFN has two bias-free matrices, so its budget is `2*d*H`. Models 2–4 allocate `4*d` trainable routing parameters and distribute `H-2` hidden units among four experts. At `H=2048`, widths are 512, 512, 511, 511. Model 4 spends `3*d` on routed logits and `d` on a sigmoid gate for the always-active shared expert; the shared expert is evaluated for every non-padding token. The sigmoid weights its output but never selects it out.\n\nModel 5 uses four width-1023 experts: two active experts plus the router contain exactly `2*d*2048` parameters per token. There are no dummy parameters. Active parameter counts include router parameters and the complete common backbone (including tied embedding/output weights), and exclude inactive expert matrices. Models 2–4 have tiny route-dependent active-count differences due to integer widths; minimum and maximum are reported.\n\nTop-1 uses the selected full-softmax probability, preserving task-loss gradients into the router. Top-2 renormalizes selected probabilities. Routing is sparse, does not drop tokens, and excludes padding from dispatch and auxiliary losses. Balance loss uses detached assignment frequencies; z-loss penalizes router log-partition magnitude. The auxiliary terms are excluded from perplexity.\n",
"id": "f003648d885b"
},
{
"cell_type": "code",
"metadata": {
"tags": [
"definitions"
]
},
"source": "class MLP(nn.Module):\n def __init__(self,d,h):\n super().__init__(); self.up=nn.Linear(d,h,bias=False); self.down=nn.Linear(h,d,bias=False)\n def forward(self,x): return self.down(F.gelu(self.up(x)))\n\nclass FFN(nn.Module):\n def __init__(self,cfg):\n super().__init__(); self.variant=cfg.variant; self.capture=False; self.last=None\n if cfg.variant==1:\n self.dense=MLP(cfg.d,cfg.hidden); self.widths=[cfg.hidden]; self.topk=1\n else:\n self.topk=1 if cfg.variant in (2,4) else 2\n if cfg.variant==5:\n assert cfg.hidden%2==0\n self.widths=[cfg.hidden//2-1]*4\n else:\n q,r=divmod(cfg.hidden-2,4); self.widths=[q+(i<r) for i in range(4)]\n self.experts=nn.ModuleList([MLP(cfg.d,h) for h in self.widths])\n self.router=nn.Linear(cfg.d,3 if cfg.variant==4 else 4,bias=False)\n if cfg.variant==4: self.shared_gate=nn.Linear(cfg.d,1,bias=False)\n def forward(self,x,valid):\n if self.variant==1: return self.dense(x),x.new_zeros(()),x.new_zeros(())\n shape=x.shape; flat=x.reshape(-1,shape[-1]); positions=valid.reshape(-1).nonzero().squeeze(1)\n tokens=flat[positions]\n logits=F.linear(tokens.float(),self.router.weight.float())\n probs=logits.softmax(-1); weights,indices=probs.topk(self.topk,dim=-1)\n if self.topk>1: weights=weights/weights.sum(-1,keepdim=True)\n mixed=torch.zeros_like(tokens)\n offset=1 if self.variant==4 else 0\n for e,expert in enumerate(self.experts[offset:]):\n row,slot=torch.where(indices==e)\n if row.numel():\n values=expert(tokens[row])*weights[row,slot,None].to(tokens.dtype)\n mixed.index_add_(0,row,values.to(mixed.dtype))\n if self.variant==4:\n gate=torch.sigmoid(F.linear(tokens.float(),self.shared_gate.weight.float())).to(tokens.dtype)\n mixed=mixed+gate*self.experts[0](tokens)\n out=torch.zeros_like(flat); out.index_copy_(0,positions,mixed)\n assignments=F.one_hot(indices,num_classes=probs.shape[-1]).float().sum(1)/self.topk\n balance=probs.shape[-1]*(assignments.detach().mean(0)*probs.mean(0)).sum()\n zloss=logits.logsumexp(-1).square().mean()\n if self.capture:\n self.last=(positions.detach(),indices.detach()+offset,weights.detach())\n return out.reshape(shape),balance,zloss\n\nclass Attention(nn.Module):\n def __init__(self,cfg):\n super().__init__(); self.heads=cfg.heads; self.hd=cfg.d//cfg.heads; self.dropout=cfg.dropout\n assert self.hd%2==0\n self.qkv=nn.Linear(cfg.d,3*cfg.d,bias=False); self.out=nn.Linear(cfg.d,cfg.d,bias=False)\n angles=torch.outer(torch.arange(cfg.context),10000.0**(-torch.arange(0,self.hd,2).float()/self.hd))\n self.register_buffer('cos',angles.cos(),persistent=False); self.register_buffer('sin',angles.sin(),persistent=False)\n def rotate(self,x,positions):\n co=self.cos[positions][:,None].to(x.dtype); si=self.sin[positions][:,None].to(x.dtype)\n a,b=x[...,::2],x[...,1::2]\n return torch.stack((a*co-b*si,a*si+b*co),-1).flatten(-2)\n def forward(self,x,valid,positions,cache=None):\n b,t,d=x.shape\n q,k,v=self.qkv(x).view(b,t,3,self.heads,self.hd).permute(2,0,3,1,4).unbind(0)\n q,k=self.rotate(q,positions),self.rotate(k,positions)\n if cache is not None:\n k=torch.cat((cache[0],k),dim=-2); v=torch.cat((cache[1],v),dim=-2)\n # Training uses right padding: no real token can see future padding.\n if cache is None:\n y=F.scaled_dot_product_attention(q,k,v,is_causal=True,dropout_p=self.dropout if self.training else 0.)\n else:\n y=F.scaled_dot_product_attention(q,k,v,attn_mask=valid[:,None,None,:],dropout_p=0.)\n return self.out(y.transpose(1,2).reshape(b,t,d)),(k,v)\n\nclass Block(nn.Module):\n def __init__(self,cfg):\n super().__init__(); self.n1=nn.LayerNorm(cfg.d); self.n2=nn.LayerNorm(cfg.d)\n self.attn=Attention(cfg); self.ffn=FFN(cfg); self.drop=nn.Dropout(cfg.dropout)\n def forward(self,x,valid,positions,cache=None):\n a,new_cache=self.attn(self.n1(x),valid,positions,cache); x=x+self.drop(a)\n local_valid=valid[:,-x.shape[1]:]\n y,bal,z=self.ffn(self.n2(x),local_valid)\n return x+self.drop(y),bal,z,new_cache\n\nclass Transformer(nn.Module):\n def __init__(self,cfg):\n super().__init__(); self.cfg=cfg\n self.embedding=nn.Embedding(cfg.vocab,cfg.d,padding_idx=PAD)\n self.blocks=nn.ModuleList([Block(cfg) for _ in range(cfg.layers)])\n self.norm=nn.LayerNorm(cfg.d); self.drop=nn.Dropout(cfg.dropout)\n self.apply(self.init)\n # Common parameters have exactly the same initialization in every variant.\n for name,p in self.named_parameters():\n if '.ffn.' not in name and p.ndim>=2:\n gen=torch.Generator().manual_seed(cfg.seed+int(hashlib.sha256(name.encode()).hexdigest()[:8],16))\n with torch.no_grad(): p.normal_(0,0.02,generator=gen)\n with torch.no_grad(): self.embedding.weight[PAD].zero_()\n @staticmethod\n def init(m):\n if isinstance(m,(nn.Linear,nn.Embedding)): nn.init.normal_(m.weight,std=0.02)\n if isinstance(m,nn.Linear) and m.bias is not None: nn.init.zeros_(m.bias)\n def forward(self,ids,labels=None,valid=None,positions=None,cache=None):\n if valid is None: valid=ids.ne(PAD)\n if positions is None: positions=torch.arange(ids.shape[1],device=ids.device)[None].expand_as(ids)\n x=self.drop(self.embedding(ids)); bal=x.new_zeros(()); z=x.new_zeros(()); next_cache=[]\n for i,block in enumerate(self.blocks):\n x,b,zz,kv=block(x,valid,positions,None if cache is None else cache[i]); bal=bal+b; z=z+zz; next_cache.append(kv)\n hidden=self.norm(x)\n # Materialize vocabulary logits only at supervised positions during training/evaluation.\n if labels is not None:\n mask=labels.ne(-100)\n logits=F.linear(hidden[mask],self.embedding.weight)\n losses=F.cross_entropy(logits.float(),labels[mask],reduction='none')\n return losses,bal/self.cfg.layers,z/self.cfg.layers,mask\n return F.linear(hidden,self.embedding.weight),next_cache\n\ndef parameter_counts(cfg):\n with torch.device('meta'): model=Transformer(cfg)\n total=sum(p.numel() for p in model.parameters())\n expert=sum(p.numel() for n,p in model.named_parameters() if '.ffn.' in n and ('experts' in n or 'dense' in n))\n common=total-expert\n widths=model.blocks[0].ffn.widths\n if cfg.variant==1: low=high=2*cfg.d*widths[0]\n elif cfg.variant==4:\n low=2*cfg.d*(widths[0]+min(widths[1:])); high=2*cfg.d*(widths[0]+max(widths[1:]))\n else:\n k=model.blocks[0].ffn.topk; low=2*cfg.d*sum(sorted(widths)[:k]); high=2*cfg.d*sum(sorted(widths)[-k:])\n return {'model':cfg.variant,'total':total,'active_min':common+cfg.layers*low,'active_max':common+cfg.layers*high,'expert_widths':widths}\n",
"execution_count": null,
"outputs": [],
"id": "12794c7bec06"
},
{
"cell_type": "markdown",
"metadata": {},
"source": "## Data and reproducible exposure\nOfficial splits are retained. Normalized exact English-target overlap is audited; training rows with an English target appearing in either holdout split are removed, and validation rows sharing a test target are removed before model selection. Within-training exact triplet duplicates are removed. This prevents the same English sentence from entering both training and evaluation but cannot detect paraphrases or all noisy alignment.\n\nThe byte-level BPE tokenizer is trained on 50,000 cleaned training triplets only, preserving Vietnamese and Japanese Unicode scripts. Each triplet yields both source directions. Overlength examples and literal reserved-token collisions are counted and dropped without silently truncating references. Test BLEU uses a fixed seed-2026 sample of 1,000 retained test triplets, paired across languages and models. Perplexity covers all retained test pairs.\n\nLength-bucketed training batches and their order depend only on the seed and shared dataset. A precomputed sequence of whole batches stays within the requested token budget; the exact realized budget is printed and is identical for every architecture. Both non-padding input tokens and supervised English/EOS tokens are counted. Padding and auxiliary losses never enter perplexity.\n",
"id": "ec7e82760a3a"
},
{
"cell_type": "code",
"metadata": {
"tags": [
"definitions"
]
},
"source": "REVISION='849990daee76e0f9e2eb9965e30e34bc1909a93d'\ndef canonical(text): return unicodedata.normalize('NFC',text).strip()\n\ndef prepare_data():\n from datasets import load_dataset\n base=ROOT/'shared'/'data'; base.mkdir(exist_ok=True)\n manifest_path=base/'manifest.json'\n if manifest_path.exists(): return json.loads(manifest_path.read_text())\n raw=load_dataset(DATASET,revision=REVISION)\n test_targets={canonical(t) for t in raw['test']['en']}\n val_targets={canonical(t) for t in raw['validation']['en']}\n held=test_targets|val_targets\n clean={}; audit={}; seen=set()\n for split in ('train','validation','test'):\n indices=[]; removed_overlap=0; removed_duplicate=0\n for i,row in enumerate(raw[split]):\n target=canonical(row['en'])\n if (split=='train' and target in held) or (split=='validation' and target in test_targets):\n removed_overlap+=1; continue\n key=tuple(canonical(row[k]) for k in ('vi','ja','en'))\n if split=='train' and key in seen: removed_duplicate+=1; continue\n if split=='train': seen.add(key)\n indices.append(i)\n clean[split]=raw[split].select(indices)\n audit[split]={'original_rows':len(raw[split]),'rows_after_overlap_dedup':len(indices),'removed_target_overlap':removed_overlap,'removed_train_triplet_duplicate':removed_duplicate}\n tok=Tokenizer(models.BPE()); tok.normalizer=normalizers.NFC()\n tok.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False); tok.decoder=decoders.ByteLevel()\n trainer=trainers.BpeTrainer(vocab_size=16000,min_frequency=2,special_tokens=['[PAD]','[BOS]','[EOS]','[SEP]'],initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),show_progress=False)\n sample=clean['train'].shuffle(seed=42).select(range(min(50000,len(clean['train']))))\n def corpus():\n for start in range(0,len(sample),1000):\n batch=sample[start:start+1000]\n yield [canonical(t) for lang in ('vi','ja','en') for t in batch[lang]]\n tok.train_from_iterator(corpus(),trainer=trainer); tok.save(str(ROOT/'shared'/'tokenizer_clean.json'))\n assert tok.get_vocab_size()==16000\n for i,s in enumerate(['[PAD]','[BOS]','[EOS]','[SEP]']): assert tok.token_to_id(s)==i\n token_stats={}\n for lang in ('vi','ja','en'):\n texts=[canonical(t) for t in clean['validation'][:1000][lang]]\n enc=tok.encode_batch(texts,add_special_tokens=False)\n assert all(tok.decode(e.ids)==t for e,t in zip(enc,texts))\n token_stats[lang]={'mean_tokens':float(np.mean([len(e.ids) for e in enc])),'roundtrip_checked':len(texts)}\n for split,ds in clean.items():\n seqs=[]; seps=[]; langs=[]; rows=[]; raw_ids=[]; dropped=Counter(); retained_rows=[]\n for start in range(0,len(ds),2000):\n batch=ds[start:start+2000]\n enc={lang:tok.encode_batch([canonical(t) for t in batch[lang]],add_special_tokens=False) for lang in ('vi','ja','en')}\n for j in range(len(batch['en'])):\n pair=[]\n for lang_id,lang in enumerate(LANGS):\n src,tgt=enc[lang][j].ids,enc['en'][j].ids\n if set(src+tgt)&{PAD,BOS,EOS,SEP}: dropped['reserved_token']+=1; continue\n ids=[BOS,*src,SEP,*tgt,EOS]\n if len(ids)-1>256: dropped['overlength_'+lang]+=1; continue\n if not src or not tgt: dropped['empty']+=1; continue\n pair.append((ids,len(src)+1,lang_id))\n # Paired filtering guarantees identical reference rows for both directions.\n if len(pair)!=2: dropped['triplets_removed']+=1; continue\n row_id=len(retained_rows)\n retained_rows.append({k:canonical(batch[k][j]) for k in ('en','vi','ja')} | {'quality_score':float(batch['quality_score'][j]),'source':batch['source'][j]})\n for ids,sep,lang_id in pair:\n seqs.append(np.asarray(ids,dtype=np.int32)); seps.append(sep); langs.append(lang_id); rows.append(row_id)\n folder=base/split; folder.mkdir(exist_ok=True)\n lengths=np.array([len(s) for s in seqs],dtype=np.int32)\n offsets=np.r_[0,lengths.cumsum(dtype=np.int64)]\n for name,arr in {'flat':np.concatenate(seqs),'offsets':offsets,'seps':np.array(seps,dtype=np.int32),'langs':np.array(langs,dtype=np.int8),'rows':np.array(rows,dtype=np.int32)}.items(): np.save(folder/(name+'.npy'),arr)\n (folder/'raw.json').write_text(json.dumps(retained_rows,ensure_ascii=False))\n audit[split].update({'retained_triplets':len(retained_rows),'pairs':len(seqs),'input_tokens':int((lengths-1).sum()),'target_tokens':int((lengths-1-np.array(seps)).sum()),'drops':dict(dropped)})\n print(split,audit[split],flush=True)\n manifest={'dataset':DATASET,'revision':REVISION,'tokenizer_sha256':sha(ROOT/'shared'/'tokenizer_clean.json'),'audit':audit,'tokenizer_stats':token_stats,'created_utc':time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime())}\n manifest_path.write_text(json.dumps(manifest,indent=2)); return manifest\n\nclass PairStore:\n def __init__(self,split):\n self.folder=ROOT/'shared'/'data'/split\n for name in ('flat','offsets','seps','langs','rows'): setattr(self,name,np.load(self.folder/(name+'.npy'),mmap_mode='r'))\n self.lengths=np.diff(self.offsets)-1\n def __len__(self): return len(self.seps)\n def get(self,i): return self.flat[self.offsets[i]:self.offsets[i+1]],int(self.seps[i]),int(self.langs[i])\n\ndef collate(ds,indices):\n length=int(max(ds.lengths[indices])); x=torch.zeros((len(indices),length),dtype=torch.long); y=torch.full_like(x,-100)\n lang=torch.tensor(np.asarray(ds.langs[indices]).copy(),dtype=torch.long); seps=torch.tensor(np.asarray(ds.seps[indices]).copy(),dtype=torch.long)\n for row,i in enumerate(indices):\n ids,sep,_=ds.get(i); ids=torch.from_numpy(np.array(ids,dtype=np.int64)); n=len(ids)-1\n x[row,:n]=ids[:-1]; y[row,sep:n]=ids[sep+1:]\n return x.to(DEVICE),y.to(DEVICE),lang.to(DEVICE),seps.to(DEVICE)\n\ndef epoch_batches(ds,cfg,epoch=0,indices=None,shuffle=True):\n rng=np.random.default_rng(cfg.seed+epoch)\n order=np.arange(len(ds)) if indices is None else np.asarray(indices).copy()\n if shuffle: rng.shuffle(order)\n batches=[]\n for start in range(0,len(order),4096):\n pool=order[start:start+4096]; pool=pool[np.argsort(ds.lengths[pool],kind='stable')]\n batch=[]; longest=0\n for i in pool:\n size=int(ds.lengths[i]); new_long=max(longest,size)\n if batch and ((len(batch)+1)*new_long>cfg.batch_tokens or len(batch)>=cfg.batch_examples):\n batches.append(np.array(batch)); batch=[]; longest=0\n batch.append(int(i)); longest=max(longest,size)\n if batch: batches.append(np.array(batch))\n if shuffle: rng.shuffle(batches)\n return batches\n\ndef schedule(ds,cfg,budget):\n result=[]; count=0; epoch=0\n while True:\n batches=epoch_batches(ds,cfg,epoch)\n assert batches\n for batch in batches:\n n=int(ds.lengths[batch].sum())\n if count+n>budget:\n assert result, 'Budget must fit at least one whole batch.'\n return result,count\n result.append(batch); count+=n\n epoch+=1\n\ndef schedule_digest(batches):\n h=hashlib.sha256()\n for b in batches: h.update(np.asarray(b,dtype='<i8').tobytes()); h.update(b'|')\n return h.hexdigest()\n\n@torch.inference_mode()\ndef evaluate_nll(model,ds,cfg,indices=None):\n model.eval(); nll=np.zeros(2); counts=np.zeros(2,dtype=np.int64); correct=np.zeros(2,dtype=np.int64)\n for batch in epoch_batches(ds,cfg,indices=indices,shuffle=False):\n x,y,lang,seps=collate(ds,batch)\n with amp(): losses,_,_,mask=model(x,labels=y)\n groups=lang[:,None].expand_as(y)[mask]\n for g in range(2):\n selected=groups==g; nll[g]+=float(losses[selected].double().sum()); counts[g]+=int(selected.sum())\n assert (counts>0).all()\n result={'nll':float(nll.sum()/counts.sum()),'ppl':math.exp(float(nll.sum()/counts.sum())),'target_tokens':int(counts.sum())}\n for g,lang in enumerate(LANGS): result.update({f'nll_{lang}':float(nll[g]/counts[g]),f'ppl_{lang}':math.exp(float(nll[g]/counts[g]))})\n return result\n",
"execution_count": null,
"outputs": [],
"id": "9e56479d709f"
},
{
"cell_type": "code",
"metadata": {
"tags": [
"definitions"
]
},
"source": "def training_run(cfg,train,val,out,budget,resume=True):\n out=Path(out); out.mkdir(exist_ok=True,parents=True)\n batches,realized=schedule(train,cfg,budget); digest=schedule_digest(batches)\n seed_all(cfg.seed); model=Transformer(cfg).to(DEVICE)\n groups=[{'params':[p for p in model.parameters() if p.ndim>=2],'weight_decay':cfg.weight_decay},\n {'params':[p for p in model.parameters() if p.ndim<2],'weight_decay':0.}]\n optimizer=torch.optim.AdamW(groups,lr=cfg.lr,betas=(.9,.95),eps=1e-8)\n scaler=torch.amp.GradScaler('cuda',enabled=DEVICE.type=='cuda',init_scale=128.)\n resume_path=out/'resume.pt'; history=[]; done=0; target_done=0; start_step=0; elapsed_before=0.\n if resume and resume_path.exists():\n saved=torch.load(resume_path,map_location='cpu',weights_only=False)\n assert saved['config']==asdict(cfg) and saved['schedule_digest']==digest\n model.load_state_dict(saved['model']); optimizer.load_state_dict(saved['optimizer']); scaler.load_state_dict(saved['scaler'])\n history=saved['history']; done=saved['input_tokens']; target_done=saved['target_tokens']; start_step=saved['step']; elapsed_before=saved['seconds']\n torch.set_rng_state(saved['torch_rng'])\n if DEVICE.type=='cuda': torch.cuda.set_rng_state_all(saved['cuda_rng'])\n del saved\n val_idx=np.sort(np.random.default_rng(109).choice(len(val),min(cfg.val_examples,len(val)),replace=False))\n started=time.monotonic(); running_nll=0.; running_targets=0; next_eval=(done//cfg.eval_interval+1)*cfg.eval_interval\n def save_resume(step):\n atomic_torch({'model':model.state_dict(),'optimizer':optimizer.state_dict(),'scaler':scaler.state_dict(),'config':asdict(cfg),\n 'schedule_digest':digest,'history':history,'input_tokens':done,'target_tokens':target_done,'step':step,\n 'seconds':elapsed_before+time.monotonic()-started,'torch_rng':torch.get_rng_state(),\n 'cuda_rng':torch.cuda.get_rng_state_all() if DEVICE.type=='cuda' else []},resume_path)\n for step in range(start_step,len(batches)):\n model.train(); indices=batches[step]; x,y,lang,seps=collate(train,indices)\n n=int(train.lengths[indices].sum()); nt=int(y.ne(-100).sum())\n progress=done/max(1,realized); warm=.03\n factor=max(.01,progress/warm) if progress<warm else .1+.9*.5*(1+math.cos(math.pi*(progress-warm)/(1-warm)))\n for group in optimizer.param_groups: group['lr']=cfg.lr*factor\n optimizer.zero_grad(set_to_none=True)\n with amp():\n losses,bal,z,_=model(x,labels=y)\n loss=losses.mean()+cfg.balance_coef*bal+cfg.z_coef*z\n if not torch.isfinite(loss): raise FloatingPointError(f'Nonfinite loss at step {step}')\n scale_before=scaler.get_scale(); scaler.scale(loss).backward(); scaler.unscale_(optimizer)\n grad_norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)\n scaler.step(optimizer); scaler.update()\n skipped=scaler.get_scale()<scale_before\n if skipped: raise FloatingPointError('AMP skipped an update; reduce loss scale and restart equally for all variants.')\n done+=n; target_done+=nt; running_nll+=float(losses.detach().sum()); running_targets+=nt\n if (step+1)%100==0 or step+1==len(batches):\n record={'kind':'train','step':step+1,'input_tokens':done,'target_tokens':target_done,'nll':running_nll/running_targets,\n 'balance':float(bal.detach()),'router_z':float(z.detach()),'grad_norm':float(grad_norm),\n 'lr':optimizer.param_groups[0]['lr'],'seconds':elapsed_before+time.monotonic()-started}\n history.append(record); running_nll=0.; running_targets=0\n print(f\"model {cfg.variant} | {done:,}/{realized:,} input tokens | nll {record['nll']:.4f} | {record['seconds']:.0f}s\",flush=True)\n if done>=next_eval or step+1==len(batches):\n score=evaluate_nll(model,val,cfg,val_idx); history.append({'kind':'validation','input_tokens':done,**score})\n print('validation',score,flush=True); next_eval+=cfg.eval_interval; save_resume(step+1)\n assert done==realized\n summary={'requested_budget':budget,'input_tokens':done,'target_tokens':target_done,'steps':len(batches),\n 'schedule_digest':digest,'seconds':elapsed_before+time.monotonic()-started,'validation_selection':history[-1]}\n atomic_torch({'model':model.state_dict(),'config':asdict(cfg),'summary':summary,'history':history,\n 'tokenizer_sha256':sha(ROOT/'shared'/'tokenizer_clean.json')},out/'checkpoint.pt')\n if budget!=cfg.token_budget and resume_path.exists(): resume_path.unlink()\n return model,summary,history\n\n@torch.inference_mode()\ndef greedy_translations(model,tok,texts,cfg,batch_size=32):\n model.eval(); results=[]; capped=[]\n encoded=[e.ids for e in tok.encode_batch([canonical(t) for t in texts],add_special_tokens=False)]\n for start in range(0,len(encoded),batch_size):\n sources=encoded[start:start+batch_size]; prompts=[[BOS,*s,SEP] for s in sources]\n assert all(len(p)<cfg.context for p in prompts),'Source exceeds context; do not silently truncate it.'\n lens=torch.tensor([len(p) for p in prompts],device=DEVICE); b=len(prompts); width=int(lens.max())\n x=torch.zeros((b,width),dtype=torch.long,device=DEVICE)\n for i,p in enumerate(prompts): x[i,:len(p)]=torch.tensor(p,device=DEVICE)\n valid=x.ne(PAD)\n with amp(): logits,cache=model(x)\n next_logits=logits[torch.arange(b,device=DEVICE),lens-1]; alive=torch.ones(b,dtype=torch.bool,device=DEVICE)\n generated=[[] for _ in prompts]; hit_cap=[False]*b\n for t in range(cfg.max_new_tokens):\n next_logits[:,[PAD,BOS,SEP]]=-float('inf'); token=next_logits.argmax(-1)\n for i in range(b):\n if alive[i]:\n if int(token[i])==EOS: alive[i]=False\n else: generated[i].append(int(token[i]))\n if alive[i] and (int(lens[i])+t>=cfg.context-1 or t==cfg.max_new_tokens-1): alive[i]=False; hit_cap[i]=True\n if not bool(alive.any()): break\n token=torch.where(alive,token,torch.full_like(token,EOS)); valid=torch.cat([valid,torch.ones((b,1),device=DEVICE,dtype=torch.bool)],1)\n positions=(lens+t).clamp(max=cfg.context-1)[:,None]\n with amp(): logits,cache=model(token[:,None],valid=valid,positions=positions,cache=cache)\n next_logits=logits[:,0]\n results.extend(tok.decode(seq) for seq in generated); capped.extend(hit_cap)\n return results,capped\n\n@torch.inference_mode()\ndef evaluate_bleu(model,tok,cfg,split='test',nrows=1000,seed=2026,bootstrap=300):\n rows=json.loads((ROOT/'shared'/'data'/split/'raw.json').read_text())\n chosen=np.sort(np.random.default_rng(seed).choice(len(rows),min(nrows,len(rows)),replace=False))\n selected=[rows[i] for i in chosen]; records=[]; result={'triplets':len(selected),'sample_seed':seed,'sample_rows_sha256':hashlib.sha256(chosen.astype('<i8').tobytes()).hexdigest()}\n metric=sacrebleu.metrics.BLEU(); started=time.monotonic()\n for lang in LANGS:\n hyps,caps=greedy_translations(model,tok,[r[lang] for r in selected],cfg)\n refs=[r['en'] for r in selected]; result['bleu_'+lang]=metric.corpus_score(hyps,[refs]).score\n result['generation_cap_'+lang]=sum(caps)\n records.extend({'row':int(i),'language':lang,'source':r[lang],'reference':ref,'hypothesis':hyp,'capped':cap,'quality_score':r['quality_score'],'corpus_source':r['source']} for i,r,ref,hyp,cap in zip(chosen,selected,refs,hyps,caps))\n result['bleu']=metric.corpus_score([r['hypothesis'] for r in records],[[r['reference'] for r in records]]).score\n result['signature']=str(metric.get_signature()); result['seconds']=time.monotonic()-started\n # Paired triplet resampling keeps vi/ja observations of the same target together.\n rng=np.random.default_rng(314); stats=[]; n=len(selected)\n for _ in range(bootstrap):\n ix=rng.integers(0,n,n); res=[records[int(i)+offset] for offset in (0,n) for i in ix]\n stats.append(metric.corpus_score([r['hypothesis'] for r in res],[[r['reference'] for r in res]]).score)\n if stats: result['bleu_ci95']=np.percentile(stats,[2.5,97.5]).tolist()\n return result,records\n",
"execution_count": null,
"outputs": [],
"id": "fff22cc41aa0"
},
{
"cell_type": "code",
"metadata": {
"tags": [
"definitions"
]
},
"source": "@torch.inference_mode()\ndef routing_diagnostics(model,ds,cfg,max_pairs=1000):\n if cfg.variant==1: return {'applicable':False,'reason':'Dense MLP has no router.'}\n for block in model.blocks: block.ffn.capture=True\n counts=np.zeros((cfg.layers,2,2,4),dtype=np.int64); totals=np.zeros((2,2),dtype=np.int64)\n indices=np.sort(np.random.default_rng(909).choice(len(ds),min(max_pairs,len(ds)),replace=False))\n for batch in epoch_batches(ds,cfg,indices=indices,shuffle=False):\n x,y,lang,seps=collate(ds,batch)\n with amp(): model(x,labels=y)\n positions=torch.arange(x.shape[1],device=DEVICE)[None].expand_as(x)\n regions=[(positions>0)&(positions<seps[:,None])&x.ne(PAD),(positions>seps[:,None])&x.ne(PAD)]\n for g in range(2):\n for r,region in enumerate(regions): totals[g,r]+=int((region&(lang[:,None]==g)).sum())\n for layer,block in enumerate(model.blocks):\n pos,expert,weights=block.ffn.last\n flat_lang=lang[:,None].expand_as(x).reshape(-1)[pos]\n for g in range(2):\n for r,region in enumerate(regions):\n mask=region.reshape(-1)[pos]&(flat_lang==g)\n if mask.any(): counts[layer,g,r]+=torch.bincount(expert[mask].flatten(),minlength=4).cpu().numpy()\n if cfg.variant==4: counts[layer,g,r,0]+=int(mask.sum())\n for block in model.blocks: block.ffn.capture=False; block.ffn.last=None\n rates=counts/np.maximum(totals[None,:,:,None],1)\n expected=1 if cfg.variant==2 else 2\n assert np.allclose(rates.sum(-1),expected)\n if cfg.variant==4: assert np.allclose(rates[...,0],1)\n diagnostics=[]\n for layer in range(cfg.layers):\n for g,lang in enumerate(LANGS):\n for r,region in enumerate(('source','English-prefix')):\n values=counts[layer,g,r,1:] if cfg.variant==4 else counts[layer,g,r]\n dist=values/max(1,values.sum()); entropy=float(-(dist[dist>0]*np.log(dist[dist>0])).sum()/np.log(len(dist)))\n diagnostics.append({'layer':layer+1,'language':lang,'region':region,'normalized_routing_entropy':entropy,'unused_experts':int((values==0).sum()),'max_assignment_share':float(dist.max())})\n return {'applicable':True,'selection_rates':rates.tolist(),'counts':counts.tolist(),'token_counts':totals.tolist(),'diagnostics':diagnostics,'sample_pairs':len(indices)}\n\ndef plot_model_results(history,routing):\n train=pd.DataFrame([r for r in history if r['kind']=='train']); val=pd.DataFrame([r for r in history if r['kind']=='validation'])\n fig,ax=plt.subplots(figsize=(8,4))\n if len(train): ax.plot(train.input_tokens/1e6,train.nll,label='Training target NLL')\n if len(val): ax.plot(val.input_tokens/1e6,val.nll,'o-',label='Validation target NLL')\n ax.set(xlabel='Non-padding training input tokens (millions)',ylabel='English target NLL'); ax.legend(); ax.grid(alpha=.2); plt.show()\n if routing.get('applicable'):\n rates=np.array(routing['selection_rates']); layers=rates.shape[0]\n fig,axes=plt.subplots(2,layers,figsize=(3*layers,6),squeeze=False,constrained_layout=True)\n for r,region in enumerate(('Source tokens','English-prefix tokens')):\n for layer in range(layers):\n ax=axes[r,layer]; im=ax.imshow(rates[layer,:,r],vmin=0,vmax=1,cmap='viridis',aspect='auto')\n ax.set(title=f'{region}, layer {layer+1}',xticks=range(4),xticklabels=[f'E{i}' for i in range(4)],yticks=[0,1],yticklabels=['vi','ja'])\n for g in range(2):\n for e in range(4): ax.text(e,g,f'{rates[layer,g,r,e]:.2f}',ha='center',va='center',color='white' if rates[layer,g,r,e]<.5 else 'black',fontsize=8)\n fig.colorbar(im,ax=axes.ravel().tolist(),label='Fraction of region tokens selecting expert'); plt.show()\n display(pd.DataFrame(routing['diagnostics']))\n\ndef run_checks(variant):\n cfg=Config(variant=variant,vocab=96,d=32,heads=4,layers=2,hidden=128,context=32,dropout=0.,profile='checks')\n seed_all(7); model=Transformer(cfg).to(DEVICE).eval()\n x=torch.tensor([[BOS,9,10,SEP,11,12,EOS],[BOS,13,SEP,14,EOS,PAD,PAD]],device=DEVICE)\n labels=torch.tensor([[-100,-100,-100,11,12,EOS,-100],[-100,-100,14,EOS,-100,-100,-100]],device=DEVICE)\n with torch.no_grad():\n first,_=model(x); alt=x.clone(); alt[0,4]=17; changed,_=model(alt)\n assert torch.allclose(first[0,:4],changed[0,:4],atol=1e-6), 'Causal leakage'\n solo,_=model(x[1:2,:5]); assert torch.allclose(solo[0],first[1,:5],atol=1e-5),'Padding changes real logits'\n prefix=x[:1,:5]; full,_=model(prefix); pref,cache=model(prefix[:,:4])\n incremental,_=model(prefix[:,4:],cache=cache,valid=torch.ones((1,5),dtype=torch.bool,device=DEVICE),positions=torch.tensor([[4]],device=DEVICE))\n assert torch.allclose(full[:,-1],incremental[:,0],atol=1e-5), 'KV cache mismatch'\n losses,bal,z,mask=model(x,labels=labels)\n direct=F.cross_entropy(first[mask].float(),labels[mask],reduction='none')\n assert torch.allclose(losses,direct,atol=1e-6)\n (losses.mean()+.01*bal+.001*z).backward()\n if variant>1:\n for block in model.blocks:\n grad=block.ffn.router.weight.grad\n assert grad is not None and torch.isfinite(grad).all() and grad.abs().sum()>0,'Router has no learning signal'\n # Sparse dispatch equals an independently assembled dense expert mixture.\n f=model.blocks[0].ffn; tokens=torch.randn((2,5,cfg.d),device=DEVICE); valid=torch.ones((2,5),device=DEVICE,dtype=torch.bool)\n got,_,_=f(tokens,valid); flat=tokens.reshape(-1,cfg.d)\n probs=f.router(flat).softmax(-1); w,idx=probs.topk(f.topk,-1)\n if f.topk>1: w=w/w.sum(-1,keepdim=True)\n ref=torch.zeros_like(flat); offset=int(variant==4)\n for e,expert in enumerate(f.experts[offset:]):\n coefficient=(w*(idx==e)).sum(-1,keepdim=True); ref+=coefficient*expert(flat)\n if variant==4: ref+=f.shared_gate(flat).sigmoid()*f.experts[0](flat)\n assert torch.allclose(got.reshape_as(ref),ref,atol=1e-6),'Sparse dispatch mismatch'\n optimizer=torch.optim.AdamW(model.parameters(),lr=.006)\n initial=float(losses.mean().detach())\n for _ in range(35):\n optimizer.zero_grad(); l,b,z,_=model(x,labels=labels); (l.mean()+.01*b+.001*z).backward(); optimizer.step()\n final=float(model(x,labels=labels)[0].mean().detach()); assert final<initial*.6,'Tiny-batch overfit failed'\n stream=io.BytesIO(); torch.save(model.state_dict(),stream); stream.seek(0)\n copy=Transformer(cfg).to(DEVICE).eval(); copy.load_state_dict(torch.load(stream,map_location=DEVICE,weights_only=True))\n with torch.no_grad(): assert torch.equal(model(x)[0],copy(x)[0]),'Checkpoint round trip failed'\n result={'model':variant,'causal_attention':'PASS','padding':'PASS','target_loss':'PASS','KV_cache':'PASS','checkpoint':'PASS','tiny_overfit_initial':initial,'tiny_overfit_final':final,'sparse_dispatch':'PASS' if variant>1 else 'not applicable','router_gradient':'PASS' if variant>1 else 'not applicable'}\n del model,copy,optimizer; gc.collect()\n if DEVICE.type=='cuda': torch.cuda.empty_cache()\n return result\n",
"execution_count": null,
"outputs": [],
"id": "20f287d7e94d"
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
} |