File size: 7,010 Bytes
9287d39 | 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 | """Shared production operations; CPU-heavy validation/export is invoked on mini."""
import hashlib,json,os,random
from pathlib import Path
import numpy as np
import torch
from dense import Config,DenseLM,prepare_layout
ROOT=Path(__file__).resolve().parent
OUT=Path('/mnt/experiments/DENSE_VALUE_DIRECT_20260911')
DATA=Path('/home/youthk/Documents/Codex/2026-08-29/new-chat-6/DST_CYCLIC_ENRICHED_100M_20260908/data/cycle')
TOKENIZER=Path('/home/youthk/デスクトップ/訓練用_共通ファイル/tokenizers/babylm_2026_strict_small_16k_bpe_98dfab9e/tokenizer.json')
ARMS=('mha_gated_ffn_router',)
SOURCE_NAMES=('dense.py','configuration_dense.py','modeling_dense.py','runtime.py','train.py','exclusive.py','test_cpu.py','gpu_acceptance.py','resume_acceptance.py','parent_dense.py')
def source_identity():return {n:sha(ROOT/n) for n in SOURCE_NAMES}
def sha(p):
h=hashlib.sha256()
with Path(p).open('rb') as f:
for block in iter(lambda:f.read(8<<20),b''):h.update(block)
return h.hexdigest()
def write_json(p,value):
p=Path(p);p.parent.mkdir(parents=True,exist_ok=True)
temp=p.with_name(p.name+'.partial');temp.write_text(json.dumps(value,indent=2)+'\n');os.replace(temp,p)
def save_checkpoint(p,value):
p=Path(p);p.parent.mkdir(parents=True,exist_ok=True)
temp=p.with_name(p.name+'.partial')
with temp.open('wb') as f:torch.save(value,f);f.flush();os.fsync(f.fileno())
os.replace(temp,p)
write_json(str(p)+'.integrity.json',dict(bytes=p.stat().st_size,sha256=sha(p)))
def setup_gpu():
torch.set_num_threads(2)
torch.manual_seed(20260907);random.seed(20260907);np.random.seed(20260907)
torch.cuda.manual_seed_all(20260907)
torch.backends.cuda.matmul.allow_tf32=True
torch.set_float32_matmul_precision('high')
torch.use_deterministic_algorithms(True)
torch.cuda.set_per_process_memory_fraction(21.5*2**30/torch.cuda.get_device_properties(0).total_memory)
def build(arm,compiled=False,value_variance=True,tensor_fusion=True):
model=DenseLM(Config(arm=arm,backend='flash',value_variance=value_variance)).cuda()
opt=torch.optim.AdamW([{'params':[p for p in model.parameters() if p.ndim>=2],'weight_decay':.1},{'params':[p for p in model.parameters() if p.ndim<2],'weight_decay':0.}],lr=8e-4,betas=(.9,.95),eps=1e-8,fused=True)
if compiled:
torch._dynamo.config.recompile_limit=64
# Real batches vary in document count. Compile stable tensor-only regions;
# keep varlen FlashAttention and its CPU-built metadata outside Dynamo.
for block in model.blocks:
block.ffn.forward=torch.compile(block.ffn.forward,fullgraph=True,dynamic=False,mode='default')
if tensor_fusion:
block.project_attention=torch.compile(block.project_attention,fullgraph=True,dynamic=False,mode='default')
block.finish_attention=torch.compile(block.finish_attention,fullgraph=True,dynamic=False,mode='default')
else:
block.attn_gate.forward=torch.compile(block.attn_gate.forward,fullgraph=True,dynamic=False,mode='default')
return model,opt
class Epoch:
def __init__(self,number,root=DATA):
self.row=json.loads((root/'dataset_manifest.json').read_text())['epochs'][number]
self.length=self.row['sequence_length'];self.rows=self.row['rows']
path=root/f"epoch_{self.row['epoch']:02d}_len_{self.length}"
self.tokens=np.memmap(path/'tokens.bin',mode='r',dtype='<u2',shape=(self.rows,self.length))
self.segments=np.memmap(path/'segments.bin',mode='r',dtype='<u2',shape=(self.rows,self.length))
self.lengths=np.memmap(path/'lengths.bin',mode='r',dtype='<u2')
self.words=np.memmap(path/'words.bin',mode='r',dtype='<u4')
def host(self,start,stop):
ids=torch.from_numpy(np.array(self.tokens[start:stop],dtype=np.int64))
seg=torch.from_numpy(np.array(self.segments[start:stop],dtype=np.int32))
valid=torch.arange(self.length)[None,:]<torch.from_numpy(np.array(self.lengths[start:stop],dtype=np.int64))[:,None]
seg=torch.where(valid,seg+1,0)
labels=torch.full_like(ids,-100)
same=(seg[:,:-1]>0)&(seg[:,:-1]==seg[:,1:])
labels[:,:-1]=torch.where(same,ids[:,1:],-100)
return ids,seg,labels,int(same.sum()),int(self.words[start:stop].sum())
def update(model,opt,data,start,stop,micro,words):
from dataclasses import replace
probability=0. if words<30_000_000 else .02 if words<60_000_000 else .05
for block in model.blocks:
if block.ffn.c.dropout!=probability:
block.ffn.c=replace(block.ffn.c,dropout=probability)
batches=[data.host(i,min(i+micro,stop)) for i in range(start,stop,micro)]
count=sum(b[3] for b in batches);exposure=sum(b[4] for b in batches)
assert count>0
opt.zero_grad(set_to_none=True);total=torch.zeros((),device='cuda')
for ids,seg,labels,n,_ in batches:
layout=prepare_layout(seg,'cuda','flash')
with torch.autocast('cuda',dtype=torch.bfloat16):loss=model(ids.cuda(),layout,labels.cuda())*(n/count)
loss.backward();total+=loss.detach()
norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True)
base_lr=8e-4 if words<30_000_000 else 7e-4 if words<60_000_000 else 6e-4
lr=base_lr*min(1.,max((words+exposure)/1_600_000,1e-4))
for group in opt.param_groups:group['lr']=lr
opt.step()
assert torch.isfinite(total)
return dict(loss=float(total),grad_norm=float(norm),lr=lr,words=exposure,targets=count,ffn_dropout=probability)
def rng():return dict(torch=torch.get_rng_state(),cuda=torch.cuda.get_rng_state_all(),numpy=np.random.get_state(),python=random.getstate())
def restore_rng(r):
torch.set_rng_state(r['torch']);torch.cuda.set_rng_state_all(r['cuda']);np.random.set_state(r['numpy']);random.setstate(r['python'])
class TailAverage:
"""Arithmetic mean of raw post-update weights, separate from training weights."""
def __init__(self):self.count=0;self.mean={};self.steps=[]
@torch.no_grad()
def add(self,model,step):
self.count+=1;self.steps.append(step)
for name,value in model.state_dict().items():
if self.count==1:self.mean[name]=value.detach().clone()
elif value.is_floating_point():self.mean[name].add_((value-self.mean[name])/self.count)
else:self.mean[name].copy_(value)
def state_dict(self):return dict(count=self.count,mean=self.mean,steps=list(self.steps))
def load_state_dict(self,state,device):
self.count=state['count'];self.steps=list(state['steps']);self.mean={k:v.to(device).clone() for k,v in state['mean'].items()}
def validate_acceptance(contract):
for name in ('CPU_ACCEPTANCE.json','GPU_ACCEPTANCE.json'):
path=ROOT/name
assert sha(path)==contract['acceptance_sha256'][name],name
report=json.loads(path.read_text())
assert report['status']=='PASS',name
assert report['source_sha256']==contract['source_sha256'],name
|