Blame for code/train.py

"""Qwen3 72M scratch training with true two-GPU DDP and complete resume state."""
from __future__ import annotations
import argparse
import bisect
from contextlib import nullcontext
from datetime import timedelta
import json
import math
import os
import random
import shutil
import signal
import time
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import Qwen3ForCausalLM, AutoTokenizer
from common import SEED, CONTEXT, VOCAB, dump, sha, now, model_config

class PackedData:
    def __init__(self, root, manifest, split):
        self.arrays = [np.memmap(Path(root)/f['name'], mode='r', dtype='<u2')
                       for f in manifest['splits'][split]['files']]
        self.ends = np.cumsum([len(a) for a in self.arrays]).tolist()
        self.blocks = (self.ends[-1]-1)//CONTEXT
        assert self.blocks > 0
    def get(self, index):
        offset = index*CONTEXT
        shard = bisect.bisect_right(self.ends, offset)
        local = offset-(self.ends[shard-1] if shard else 0)
        a = self.arrays[shard][local:local+CONTEXT+1]
        if len(a) < CONTEXT+1:
            a = np.concatenate([a, self.arrays[shard+1][:(CONTEXT+1-len(a))]])
        return np.asarray(a, dtype=np.int64)
    def batch(self, indices, device):
        a = np.stack([self.get(i) for i in indices])
        t = torch.from_numpy(a).to(device)
        return t[:,:-1].contiguous(), t[:,1:].contiguous()

def permutation(index, count, seed=SEED):
    epoch, pos = divmod(index, count)
    rng = random.Random(seed+epoch)
    a = rng.randrange(1,max(2,count))
    while math.gcd(a,count) != 1: a = (a+1)%count or 1
    b = rng.randrange(count)
    return (a*pos+b)%count

def indices_for(step, micro, batch_size, accumulation, world, rank, count):
    start = ((step*accumulation+micro)*world+rank)*batch_size
    return [permutation(start+i,count) for i in range(batch_size)]

def lr_for(step, total, warmup, peak, minimum):
    if step < warmup: return peak*(step+1)/max(1,warmup)
    progress = min(1.0,(step-warmup)/max(1,total-warmup))
    return minimum+(peak-minimum)*0.5*(1+math.cos(math.pi*progress))

def rng_state(device):
    state = {'python':random.getstate(), 'numpy':np.random.get_state(),
             'torch':torch.get_rng_state()}
    if device.type == 'cuda': state['cuda'] = torch.cuda.get_rng_state(device)
    return state

def restore_rng(state,device):
    random.setstate(state['python']); np.random.set_state(state['numpy'])
    torch.set_rng_state(state['torch'].cpu())
    if device.type == 'cuda': torch.cuda.set_rng_state(state['cuda'].cpu(),device)

def save_checkpoint(out, model, optimizer, scaler, step, best, settings,
                    rank, world, device, best_now=False):
    states = [None]*world if rank == 0 else None
    local = rng_state(device)
    if world > 1: dist.gather_object(local, states, dst=0)
    else: states = [local]
    if rank == 0:
        directory = out/'checkpoints'; directory.mkdir(parents=True,exist_ok=True)
        path = directory/f'step-{step:08d}.pt'
        payload = {'format_version':1, 'step':step, 'best_validation_loss':best,
                   'model':model.state_dict(), 'optimizer':optimizer.state_dict(),
                   'scaler':scaler.state_dict(), 'rng_by_rank':states,
                   'settings':settings, 'data_cursor_global_blocks':step*settings['blocks_per_step'],
                   'scheduler':{'type':'warmup_cosine','next_step':step},
                   'created_utc':now()}
        tmp = path.with_suffix('.pt.partial'); torch.save(payload,tmp); tmp.replace(path)
        info = {'path':path.name,'step':step,'sha256':sha(path),'bytes':path.stat().st_size,
                'corpus_manifest_sha256':settings['corpus_manifest_sha256']}
        dump(directory/'latest.json',info)
        if best_now:
            best_path = directory/'best.pt'
            shutil.copyfile(path,best_path.with_suffix('.pt.partial'))
            best_path.with_suffix('.pt.partial').replace(best_path)
            dump(directory/'best.json',{**info,'path':'best.pt','validation_loss':best})
        # Delete only generated checkpoint files in this run's checked directory.
        checked_root = directory.resolve()
        for old in sorted(directory.glob('step-*.pt'))[:-2]:
            assert old.resolve().parent == checked_root
            old.unlink()
        print('CHECKPOINT_SAVED',json.dumps(info),flush=True)
        if os.environ.get('THINK72_CLOUD_RESUME')=='1':
            try:
                from checkpoint_cloud import backup_checkpoint
                backup_checkpoint(out,path)
            except Exception as exc:
                # Retain local checkpoints and log only exception type, never secrets/URLs.
                print('CLOUD_BACKUP_RETRY_AT_NEXT_CHECKPOINT',type(exc).__name__,flush=True)
    if world > 1: dist.barrier()

@torch.no_grad()
def evaluate(model, data, device, rank, world, batches, batch_size, amp):
    model.eval()
    totals = torch.zeros(2,device=device,dtype=torch.float64)
    for i in range(batches):
        indices = [permutation((i*world+rank)*batch_size+j,data.blocks,SEED+71)
                   for j in range(batch_size)]
        x,y = data.batch(indices,device)
        with torch.autocast(device_type=device.type,dtype=torch.float16,enabled=amp):
            logits = model(input_ids=x,use_cache=False).logits
            loss = F.cross_entropy(logits.reshape(-1,VOCAB).float(),y.reshape(-1),
                                   ignore_index=0,reduction='sum')
        totals[0] += loss.double(); totals[1] += (y != 0).sum()
    if world > 1: dist.all_reduce(totals)
    model.train()
    nll = (totals[0]/totals[1]).item()
    return {'loss':nll,'perplexity':math.exp(min(50,nll)),'evaluated_target_tokens':int(totals[1].item())}

@torch.no_grad()
def generate_samples(model, tokenizer, device, amp, max_new_tokens=96):
    model.eval()
    generations=[]
    for prompt in ['What is 12 plus 7?', 'If a box has 3 rows of 4 apples, how many apples are there?',
                   'Write a Python function that returns the square of a number.']:
        inputs=tokenizer.apply_chat_template([{'role':'user','content':prompt}],
                     add_generation_prompt=True,return_tensors='pt',return_dict=True,
                     tokenizer_kwargs={'return_token_type_ids':False}).to(device)
        with torch.autocast(device_type=device.type,dtype=torch.float16,enabled=amp):
            generated=model.generate(**inputs,max_new_tokens=max_new_tokens,do_sample=False,use_cache=True,
                                     pad_token_id=0,eos_token_id=2)
        generations.append({'prompt':prompt,'response':tokenizer.decode(generated[0,inputs['input_ids'].shape[1]:],
                                             skip_special_tokens=False,clean_up_tokenization_spaces=False)})
    return generations

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--data',required=True); ap.add_argument('--out',required=True)
    ap.add_argument('--resume'); ap.add_argument('--device',default='cuda',choices=['cpu','cuda'])
    ap.add_argument('--require-two-t4',action='store_true')
    ap.add_argument('--batch-size',type=int,default=8); ap.add_argument('--grad-accum',type=int,default=4)
    ap.add_argument('--target-tokens',type=int,default=1_500_000_000)
    ap.add_argument('--warmup-steps',type=int,default=150)
    ap.add_argument('--lr',type=float,default=6e-4); ap.add_argument('--min-lr',type=float,default=6e-5)
    ap.add_argument('--save-every',type=int,default=250); ap.add_argument('--save-seconds',type=int,default=600)
    ap.add_argument('--eval-every',type=int,default=250); ap.add_argument('--eval-batches',type=int,default=16)
    ap.add_argument('--max-hours',type=float,default=10.5)
    ap.add_argument('--stop-after-steps',type=int,default=0)
    ap.add_argument('--threads',type=int,default=2)
    args=ap.parse_args()
    rank=int(os.environ.get('RANK','0')); world=int(os.environ.get('WORLD_SIZE','1'))
    local_rank=int(os.environ.get('LOCAL_RANK','0'))
    torch.set_num_threads(args.threads)
    if args.require_two_t4:
        names=[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
        assert args.device=='cuda' and len(names)==2 and world==2 and all('T4' in n for n in names), (names,world)
    device=torch.device(f'cuda:{local_rank}' if args.device=='cuda' else 'cpu')
    if device.type=='cuda': torch.cuda.set_device(device)
    if world>1:
        dist.init_process_group('nccl' if device.type=='cuda' else 'gloo',timeout=timedelta(minutes=30))
    random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
    if device.type=='cuda': torch.cuda.manual_seed_all(SEED)
    out=Path(args.out); out.mkdir(parents=True,exist_ok=True)
    manifest_path=Path(args.data)/'corpus_manifest.json'
    manifest=json.loads(manifest_path.read_text())
    assert manifest['context']==CONTEXT and manifest['vocab_size']==VOCAB
    if rank==0:
        for split in manifest['splits'].values():
            for item in split['files']:
                p=Path(args.data)/item['name']
                assert p.stat().st_size==item['bytes'] and sha(p)==item['sha256'], p
        assert sha(Path(args.data)/'tokenizer'/'tokenizer.json')==manifest['tokenizer_sha256']
    if world>1: dist.barrier()
    datasets={s:PackedData(args.data,manifest,s) for s in ['train','validation','test']}
    config=model_config(); config._attn_implementation='sdpa'
    model=Qwen3ForCausalLM(config).to(device)
    assert sum(p.numel() for p in model.parameters())==72_000_000
    decay=[p for p in model.parameters() if p.ndim>=2]
    no_decay=[p for p in model.parameters() if p.ndim<2]
    optimizer=torch.optim.AdamW([{'params':decay,'weight_decay':0.1},
                                {'params':no_decay,'weight_decay':0.0}],
                               lr=args.lr,betas=(0.9,0.95),eps=1e-8,
                               fused=device.type=='cuda')
    amp=device.type=='cuda'
    scaler=torch.amp.GradScaler('cuda',enabled=amp)
    blocks_per_step=args.batch_size*args.grad_accum*world
    tokens_per_step=blocks_per_step*CONTEXT
    total_steps=math.ceil(args.target_tokens/tokens_per_step)
    settings={'parameters':72_000_000,'context':CONTEXT,'vocab':VOCAB,
              'seed':SEED,'world_size':world,'batch_size':args.batch_size,'grad_accum':args.grad_accum,
              'blocks_per_step':blocks_per_step,'tokens_per_step':tokens_per_step,
              'target_tokens':args.target_tokens,'total_steps':total_steps,
              'warmup_steps':args.warmup_steps,'lr':args.lr,'min_lr':args.min_lr,
              'optimizer':'AdamW betas=(0.9,0.95) eps=1e-8 decay=0.1 matrices only',
              'precision':'FP16 autocast, FP32 weights/optimizer, GradScaler' if amp else 'FP32',
              'corpus_manifest_sha256':sha(manifest_path),'tokenizer_sha256':manifest['tokenizer_sha256'],
              'model_config':config.to_dict()}
    # Exclude library-version metadata from equality; training configuration remains strict.
    settings['model_config'].pop('transformers_version',None)
    step=0; best=math.inf
    checkpoint=None
    if args.resume:
        checkpoint=torch.load(args.resume,map_location='cpu',weights_only=False)
        assert checkpoint['settings']==settings, 'Resume settings or dataset changed'
        model.load_state_dict(checkpoint['model']); optimizer.load_state_dict(checkpoint['optimizer'])
        scaler.load_state_dict(checkpoint['scaler']); step=checkpoint['step']; best=checkpoint['best_validation_loss']
        # Verify deserialization before any further training; GPU arithmetic after
        # restart can differ slightly without deterministic CUDA kernels.
        assert all(torch.equal(v.detach().cpu(),checkpoint['model'][k]) for k,v in model.state_dict().items())
        loaded_optimizer=optimizer.state_dict()
        for key, state in loaded_optimizer['state'].items():
            for field, value in state.items():
                expected=checkpoint['optimizer']['state'][key][field]
                if torch.is_tensor(value):assert torch.equal(value.cpu(),expected.cpu()), (key,field)
                else:assert value==expected, (key,field)
        if rank==0:print('RESUME_STATE_VERIFIED_EXACT',step,flush=True)
        assert checkpoint['data_cursor_global_blocks']==step*blocks_per_step
        assert checkpoint['scheduler']['next_step']==step
    wrapped=DDP(model,device_ids=[local_rank],broadcast_buffers=False,gradient_as_bucket_view=True) if world>1 else model
    if checkpoint:
        restore_rng(checkpoint['rng_by_rank'][rank],device)
        del checkpoint
    started=time.monotonic(); last_save=started; stopped=[False]
    def request_stop(*_): stopped[0]=True
    signal.signal(signal.SIGTERM,request_stop); signal.signal(signal.SIGINT,request_stop)
    if rank==0:
        hardware={'gpu_names':[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())],
                  'world_size':world,'torch':torch.__version__,'cuda':torch.version.cuda,
                  'parameters':72_000_000,'resumed_step':step}
        dump(out/'hardware.json',hardware); dump(out/'training_settings.json',settings)
        print('HARDWARE_VERIFIED',json.dumps(hardware),flush=True)
        dump(out/'run_status.json',{'status':'RUNNING','step':step,'total_steps':total_steps,'started_utc':now()})
    initial=evaluate(model,datasets['validation'],device,rank,world,min(4,args.eval_batches),args.batch_size,amp)
    if rank==0:
        print('INITIAL_VALIDATION',json.dumps(initial),flush=True)
        dump(out/'initial_validation.json',initial)
    model.train(); start_step=step; overflows=0; latest_eval=None
    while step<total_steps:
        iter_started=time.monotonic(); optimizer.zero_grad(set_to_none=True)
        lr=lr_for(step,total_steps,args.warmup_steps,args.lr,args.min_lr)
        for group in optimizer.param_groups: group['lr']=lr
        losses=torch.zeros(1,device=device)
        for micro in range(args.grad_accum):
            ids=indices_for(step,micro,args.batch_size,args.grad_accum,world,rank,datasets['train'].blocks)
            x,y=datasets['train'].batch(ids,device)
            sync=wrapped.no_sync() if world>1 and micro<args.grad_accum-1 else nullcontext()
            with sync:
                with torch.autocast(device_type=device.type,dtype=torch.float16,enabled=amp):
                    logits=wrapped(input_ids=x,use_cache=False).logits
                    loss=F.cross_entropy(logits.reshape(-1,VOCAB).float(),y.reshape(-1),ignore_index=0)
                finite=torch.isfinite(loss.detach()).to(torch.int32)
                if world>1: dist.all_reduce(finite,op=dist.ReduceOp.MIN)
                if not finite.item(): raise FloatingPointError(f'Non-finite loss at step {step}')
                losses+=loss.detach()/args.grad_accum
                scaler.scale(loss/args.grad_accum).backward()
                del logits,loss
        scaler.unscale_(optimizer)
        grad_norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
        before_scale=scaler.get_scale()
        scaler.step(optimizer); scaler.update()
        if scaler.get_scale()<before_scale:
            overflows+=1
            if rank==0: print('AMP_RETRY',step,'scale',scaler.get_scale(),flush=True)
            if overflows>32: raise FloatingPointError('Repeated FP16 overflows')
            continue
        overflows=0; step+=1
        if world>1: dist.all_reduce(losses); losses/=world
        if step==1 or step%10==0:
            elapsed=time.monotonic()-started
            record={'step':step,'total_steps':total_steps,'tokens_seen':step*tokens_per_step,
                    'train_loss':losses.item(),'lr':lr,'grad_norm':float(grad_norm),
                    'step_seconds':time.monotonic()-iter_started,
                    'tokens_per_second':(step-start_step)*tokens_per_step/max(1,elapsed),
                    'gpu_memory_peak_gb':torch.cuda.max_memory_allocated(device)/1e9 if amp else 0}
            if rank==0:
                print('TRAIN',json.dumps(record),flush=True)
                with (out/'metrics.jsonl').open('a',encoding='utf-8') as f:f.write(json.dumps(record)+'\n')
                dump(out/'run_status.json',{'status':'RUNNING',**record,'updated_utc':now()})
        best_now=False
        if step%args.eval_every==0 or step==total_steps:
            latest_eval=evaluate(model,datasets['validation'],device,rank,world,args.eval_batches,args.batch_size,amp)
            best_now=latest_eval['loss']<best
            if best_now: best=latest_eval['loss']
            if rank==0:
                print('VALIDATION',step,json.dumps(latest_eval),flush=True)
                dump(out/'validation_latest.json',{'step':step,**latest_eval})
        stop=torch.tensor(int(stopped[0] or time.monotonic()-started>=args.max_hours*3600
                              or (args.stop_after_steps and step>=args.stop_after_steps)),device=device)
        save=torch.tensor(int(step==1 or step%args.save_every==0 or time.monotonic()-last_save>=args.save_seconds
                              or best_now or step==total_steps),device=device)
        if world>1:
            dist.all_reduce(stop,op=dist.ReduceOp.MAX); dist.all_reduce(save,op=dist.ReduceOp.MAX)
        if save.item() or stop.item():
            save_checkpoint(out,model,optimizer,scaler,step,best,settings,rank,world,device,best_now)
            last_save=time.monotonic()
        if stop.item(): break
    final_eval=evaluate(model,datasets['test'],device,rank,world,args.eval_batches,args.batch_size,amp) if step>=total_steps else None
    if rank==0:
        export=out/'model'; export.mkdir(exist_ok=True)
        model.config.use_cache=True
        model.save_pretrained(export,max_shard_size='2GB',safe_serialization=True)
        tokenizer=AutoTokenizer.from_pretrained(Path(args.data)/'tokenizer',local_files_only=True)
        tokenizer.model_input_names=['input_ids','attention_mask']
        tokenizer.init_kwargs['model_input_names']=tokenizer.model_input_names
        tokenizer.save_pretrained(export)
        generations=[]
        sample_error=None
        try:
            generations=generate_samples(model,tokenizer,device,amp)
        except Exception as exc:
            # Optional previews must not prevent final status/backup after a saved checkpoint.
            sample_error={'type':type(exc).__name__,'message':str(exc)}
            dump(out/'sample_generation_error.json',sample_error)
            print('SAMPLE_GENERATION_FAILED',type(exc).__name__,flush=True)
        dump(out/'sample_generations.json',generations)
        for name in ['source_manifest.json','corpus_manifest.json']:
            shutil.copyfile(Path(args.data)/name,out/name)
        status='COMPLETE' if step>=total_steps else 'CHECKPOINTED_NEEDS_RESUME'
        card=f'''---
language: en
library_name: transformers
pipeline_tag: text-generation
license: other
---
# Think72 Qwen3-style scratch checkpoint

Status: {status}. Completed optimizer updates: {step}/{total_steps}.
Processed token positions: {step*tokens_per_step:,}. Context: 1024. Parameters: 72,000,000.
This checkpoint uses 14 Qwen3 decoder layers, hidden size 512, intermediate size 1536,
8 query heads and 4 KV heads of dimension 128, Q/K RMSNorm, RoPE, SwiGLU and tied embeddings.
Weights and the 33,069-token byte-level BPE were trained from scratch; no pretrained Qwen weights were loaded.

Data: https://www.kaggle.com/datasets/faithdemo/think72-english-tokens-v1
Training run: https://www.kaggle.com/code/faithdemo/think72-qwen3-t4x2
Exact corpus manifest SHA-256: {settings['corpus_manifest_sha256']}.
Exact tokenizer SHA-256: {settings['tokenizer_sha256']}.
See the adjacent source/corpus manifests, training_settings.json and DATA_CARD.md for provenance,
schema, splits, loss objective, learning-rate schedule and third-party attribution.
Original source revisions are unknown; local file hashes identify the consumed snapshot.

Load with AutoTokenizer.from_pretrained(path) and AutoModelForCausalLM.from_pretrained(path).
Use tokenizer.apply_chat_template(messages, add_generation_prompt=True) and keep the total
prompt plus generated tokens at or below 1024. Optional reasoning text uses <think> and </think>.
Intended language is English; tasks include experimental text continuation, math, science and code.
Validation: {json.dumps(latest_eval)}. Held-out test sample: {json.dumps(final_eval)}.
Perplexity is not a task-accuracy score; inspect sample_generations.json and evaluate independently.

Limitations: tiny model capacity, short context, mostly reasoning/synthetic rather than broad
language pretraining, possible overlap between semantically related questions, factual errors,
repetition and unreliable reasoning. Qwen3 architecture does not confer Qwen3-0.6B capabilities.

Copyright (c) 2026 sraivante covers original contributions and selection/arrangement only,
under Apache License 2.0. Third-party source terms remain in force, including CC-BY-NC-4.0
material; no blanket Apache or commercial-use license for these weights/data is asserted.
'''
        (export/'README.md').write_text(card,encoding='utf8')
        if Path('DATA_CARD.md').exists():shutil.copy2('DATA_CARD.md',out/'DATA_CARD.md')
        result={'status':status,
                'step':step,'total_steps':total_steps,'tokens_seen':step*tokens_per_step,
                'validation':latest_eval,'heldout_test':final_eval,'ended_utc':now(),
                'sample_generation_error':sample_error,
                'model_files':{p.name:{'bytes':p.stat().st_size,'sha256':sha(p)} for p in export.iterdir() if p.is_file()},
                'limitations':'A scratch-trained 72M model with a 1024-token window; Qwen3 architecture does not imply Qwen3-0.6B reasoning ability. Task accuracy requires separate evaluation.'}
        dump(out/'run_status.json',result)
        print('TRAINING_SESSION_FINISHED',json.dumps(result),flush=True)
    if world>1: dist.barrier(); dist.destroy_process_group()

if __name__=='__main__': main()