"""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=' 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 step1 and micro1: 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()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']=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 and . 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()