training_core.py

1"""Full-game causal training primitives and stable local checkpoint loading API."""
2from contextlib import nullcontext
3import hashlib
4import json
5import math
6import os
7from pathlib import Path
8import random
9import shutil
10import importlib.util
11import sys
12import torch
13import torch.distributed as dist
14import torch.nn.functional as F
15from transformers import GPT2Config,GPT2LMHeadModel
16# Expose only this existing pure-Python package, never its old site's compiled
17# Torch/SciPy modules. This reads existing files without installing or copying.
18if importlib.util.find_spec('chess') is None:
19    chess_root=Path('<PRIVATE_PATH>')
20    spec=importlib.util.spec_from_file_location('chess',chess_root/'__init__.py',submodule_search_locations=[str(chess_root)])
21    chess_module=importlib.util.module_from_spec(spec);sys.modules['chess']=chess_module;spec.loader.exec_module(chess_module)
22from chessdata import BOS,EOS,PAD,VERSION
23
24CHECKPOINT_FORMAT='causal-chess-train-v1'
25BASELINE=dict(vocab_size=128,n_positions=1024,n_ctx=1024,n_embd=768,n_layer=12,n_head=12,
26              resid_pdrop=0.,embd_pdrop=0.,attn_pdrop=0.,use_cache=False,
27              bos_token_id=BOS,eos_token_id=EOS,pad_token_id=PAD,attn_implementation='sdpa')
28
29
30def digest_json(value):
31    return hashlib.sha256(json.dumps(value,sort_keys=True,separators=(',',':')).encode()).hexdigest()
32
33
34def model(config=None,device='cpu'):
35    config=BASELINE if config is None else config
36    cfg=GPT2Config(**config)
37    cfg._attn_implementation='sdpa'
38    return GPT2LMHeadModel(cfg).to(device)
39
40
41def collate(samples, device='cpu'):
42    """Inputs/labels are ALREADY shifted. Never pass labels to HF forward."""
43    if not samples:
44        return (torch.tensor([[PAD]],device=device),torch.tensor([[-100]],device=device),
45                torch.ones((1,1),dtype=torch.bool,device=device))
46    lengths=[len(s)-1 for s in samples]
47    if any(n<1 or n>1024 or s[0]!=BOS or s[-1]!=EOS for n,s in zip(lengths,samples)):
48        raise ValueError('full game required; no prefix crop or synthetic EOS')
49    x=torch.full((len(samples),max(lengths)),PAD,dtype=torch.long,device=device)
50    y=torch.full_like(x,-100); mask=torch.zeros_like(x,dtype=torch.bool)
51    for i,(sample,n) in enumerate(zip(samples,lengths)):
52        if isinstance(sample,bytes):sample=list(sample)  # expand only this batch row
53        x[i,:n]=torch.tensor(sample[:-1],device=device);y[i,:n]=torch.tensor(sample[1:],device=device);mask[i,:n]=True
54    return x,y,mask
55
56
57def batches(lengths,epoch,seed,micro,world,bucket=256):
58    """Every game exactly once globally, no distributed sampler duplication."""
59    rng=random.Random(seed+epoch); indices=list(range(len(lengths)));rng.shuffle(indices)
60    result=[]; size=micro*world
61    for start in range(0,len(indices),bucket):
62        ordered=sorted(indices[start:start+bucket],key=lambda i:lengths[i])
63        result.extend(ordered[i:i+size] for i in range(0,len(ordered),size))
64    rng.shuffle(result)
65    return result
66
67
68def groups(plan,accum):
69    return [plan[i:i+accum] for i in range(0,len(plan),accum)]
70
71
72def optimizer(net,lr=3e-4,fused=False):
73    decay=[];nodecay=[]
74    for p in net.parameters():
75        (decay if p.ndim>=2 else nodecay).append(p)
76    return torch.optim.AdamW([{'params':decay,'weight_decay':0.1},{'params':nodecay,'weight_decay':0.0}],
77                             lr=lr,betas=(0.9,0.95),eps=1e-8,fused=fused)
78
79
80def scheduler(opt,total,warmup):
81    def scale(step):
82        if step<warmup:return (step+1)/max(1,warmup)
83        progress=min(1.,max(0.,(step-warmup)/max(1,total-warmup)))
84        return 0.1+0.9*0.5*(1+math.cos(math.pi*progress))
85    return torch.optim.lr_scheduler.LambdaLR(opt,scale)
86
87
88def backward_group(net,samples,group,micro,rank=0,world=1,device='cpu',bf16=False):
89    """DDP averages gradients, hence world/global-valid-targets scale per sum.
90
91    All ranks execute equal microstep counts, including zero-target dummy ranks.
92    Each accumulation group denominator includes every real target on every rank.
93    """
94    denominator=sum(len(samples[i])-1 for batch in group for i in batch)
95    if denominator<=0:raise ValueError('empty update')
96    loss_sum=torch.zeros((),device=device);padded=0;valid=0
97    for number,batch in enumerate(group):
98        ids=batch[rank*micro:(rank+1)*micro]
99        x,y,mask=collate([samples[i] for i in ids],device)
100        padded+=x.numel();valid+=int((y!=-100).sum())
101        sync=net.no_sync() if hasattr(net,'no_sync') and number<len(group)-1 else nullcontext()
102        with sync:
103            with torch.autocast('cuda',dtype=torch.bfloat16,enabled=bf16):
104                logits=net(input_ids=x,attention_mask=mask,use_cache=False).logits
105                loss=F.cross_entropy(logits.float().reshape(-1,128),y.reshape(-1),ignore_index=-100,reduction='sum')
106            (loss*(world/denominator)).backward()
107        loss_sum+=loss.detach()
108    stats=torch.tensor([float(loss_sum),valid,padded],dtype=torch.float64,device=device)
109    if world>1:dist.all_reduce(stats)
110    if int(stats[1])!=denominator:raise RuntimeError('target denominator mismatch across ranks')
111    return {'loss_sum':stats[0].item(),'valid_targets':int(stats[1]),'padded_tokens':int(stats[2])}
112
113
114def rng_state(device):
115    return {'python':random.getstate(),'torch':torch.get_rng_state(),
116            'cuda':torch.cuda.get_rng_state(device) if str(device).startswith('cuda') else None}
117
118
119def restore_rng(state,device):
120    random.setstate(state['python']);torch.set_rng_state(state['torch'])
121    if state['cuda'] is not None:torch.cuda.set_rng_state(state['cuda'],device)
122
123
124def atomic_save(path,payload):
125    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True);tmp=path.with_name(path.name+'.tmp')
126    with tmp.open('wb') as stream:
127        torch.save(payload,stream);stream.flush();os.fsync(stream.fileno())
128    os.replace(tmp,path)
129    fd=os.open(path.parent,os.O_RDONLY|os.O_DIRECTORY)
130    try:os.fsync(fd)
131    finally:os.close(fd)
132
133
134def link_best(last,best):
135    last,best=Path(last),Path(best);tmp=best.with_name(best.name+'.tmp')
136    if tmp.exists():tmp.unlink()
137    try:os.link(last,tmp)
138    except OSError:shutil.copyfile(last,tmp)
139    os.replace(tmp,best)
140
141
142def checkpoint_payload(net,opt,sched,identity,cursor,rngs,metrics):
143    raw=net.module if hasattr(net,'module') else net
144    return {'format':CHECKPOINT_FORMAT,'model_config':raw.config.to_dict(),'model_state_dict':raw.state_dict(),
145            'optimizer_state_dict':opt.state_dict(),'scheduler_state_dict':sched.state_dict(),
146            'identity':identity,'identity_sha256':digest_json(identity),'cursor':cursor,'rng_by_rank':rngs,'metrics':metrics}
147
148
149def load_checkpoint(path,identity=None):
150    saved=torch.load(path,map_location='cpu',weights_only=True)
151    if (saved.get('format')!=CHECKPOINT_FORMAT or digest_json(saved['identity'])!=saved['identity_sha256']
152            or (identity is not None and saved['identity']!=identity)):
153        raise ValueError('incompatible or inconsistent checkpoint identity')
154    if saved['identity']['tokenizer']['version']!=VERSION:raise ValueError('unsupported tokenizer contract')
155    return saved
156
157
158def load_model_checkpoint(path,device='cpu'):
159    """Harness contract: return (eval GPT2LMHeadModel, checkpoint metadata).
160
161    Uses saved config/state, never downloads. Caller owns legal masking. Raw logits
162    remain 128 classes. Only load locally trusted project checkpoints.
163    """
164    saved=load_checkpoint(path)
165    net=model(saved['model_config'],device)
166    net.load_state_dict(saved['model_state_dict'],strict=True);net.eval()
167    metadata={k:saved[k] for k in ('format','model_config','identity','identity_sha256','cursor','metrics')}
168    return net,metadata
169
170
171def launch_owned_workers(script,arguments):
172    """Two local DDP workers stay in our guardian-owned process group.
173
174    Unlike torchrun's elastic subprocess launcher, never create worker sessions.
175    FileStore rendezvous is local to /tmp; no port selection race or new service.
176    """
177    import signal
178    import subprocess
179    import tempfile
180    import time
181    if os.name!='posix':raise RuntimeError('owned DDP launcher requires Linux')
182    workers=[];stopped=[None]
183    def on_signal(signum,frame):stopped[0]=signum
184    previous={s:signal.signal(s,on_signal) for s in (signal.SIGTERM,signal.SIGINT)}
185    try:
186        with tempfile.TemporaryDirectory(prefix='chess-ddp-') as folder:
187            try:
188                for rank in (0,1):
189                    if stopped[0] is not None:return 128+stopped[0]
190                    env=dict(os.environ,RANK=str(rank),LOCAL_RANK=str(rank),WORLD_SIZE='2',LOCAL_WORLD_SIZE='2',
191                             CHESS_DDP_INIT_FILE=str(Path(folder)/'rendezvous'))
192                    workers.append(subprocess.Popen([sys.executable,str(script),*arguments],env=env,start_new_session=False))
193                while True:
194                    if stopped[0] is not None:return 128+stopped[0]
195                    codes=[p.poll() for p in workers]
196                    failure=next((code for code in codes if code is not None and code!=0),None)
197                    if failure is not None:return failure if failure>0 else 128-failure
198                    if all(code==0 for code in codes):return 0
199                    time.sleep(.05)
200            finally:
201                for p in workers:
202                    if p.poll() is None:p.terminate()
203                until=time.monotonic()+2
204                while any(p.poll() is None for p in workers) and time.monotonic()<until:time.sleep(.02)
205                for p in workers:
206                    if p.poll() is None:p.kill()
207                for p in workers:p.wait(timeout=5)
208    finally:
209        for signum,handler in previous.items():signal.signal(signum,handler)
210