#!/usr/bin/env python """QLoRA trainer on the Gemma-4-12B QAT foundation. 4-bit NF4 base + Liger fused CE, gradient checkpointing OFF (memory freed by Liger). Assistant-only masking by incremental re-render. Cache keyed on CORPUS hash (handover ยง4 footgun).""" import os,sys,json,math,time,argparse,hashlib,random import torch os.environ.setdefault("HF_HOME","/workspace/.hf_home") from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from peft import LoraConfig, get_peft_model def encode(tok, msgs): """Supervise ONLY assistant spans, found by incremental re-render - never regex.""" ids=[]; lab=[] for i,m in enumerate(msgs): if m["role"]!="assistant": continue p=tok.apply_chat_template(msgs[:i], add_generation_prompt=True, tokenize=False) f=tok.apply_chat_template(msgs[:i+1], tokenize=False) if not f.startswith(p): p=tok.apply_chat_template(msgs[:i], tokenize=False) if not f.startswith(p): return None pi=tok(p, add_special_tokens=False)["input_ids"] fi=tok(f, add_special_tokens=False)["input_ids"] if len(fi)<=len(pi): return None ids=fi; lab=[-100]*len(pi)+fi[len(pi):] if not ids or all(l==-100 for l in lab): return None return ids,lab def build_blocks(tok, path, seq, cache): if os.path.exists(cache): d=torch.load(cache); print(f"[data] {len(d['ids'])} blocks from cache",flush=True) return d["ids"],d["lab"] S=[];L=[] n=0; skipped=0 with open(path,encoding="utf-8") as fh: for line in fh: r=json.loads(line); e=encode(tok,r["messages"]) if e is None: skipped+=1; continue i,l=e; S+=i+[tok.eos_token_id]; L+=l+[tok.eos_token_id]; n+=1 if n%2000==0: print(f" encoded {n}",flush=True) nb=len(S)//seq ids=torch.tensor([S[i*seq:(i+1)*seq] for i in range(nb)],dtype=torch.long) lab=torch.tensor([L[i*seq:(i+1)*seq] for i in range(nb)],dtype=torch.long) keep=(lab!=-100).any(dim=1) # blocks with zero supervision -> NaN eval, no grad ids,lab=ids[keep],lab[keep] print(f"[data] {n} recs ({skipped} skipped) -> {len(ids)} blocks x {seq} " f"= {len(ids)*seq/1e6:.2f}M packed, {(lab!=-100).sum().item()/1e6:.2f}M supervised",flush=True) torch.save({"ids":ids,"lab":lab},cache) return ids,lab def main(): ap=argparse.ArgumentParser() ap.add_argument("--corpus",required=True); ap.add_argument("--out",required=True) ap.add_argument("--model",default="/workspace/models/gemma4-qat-unquantized") ap.add_argument("--seq",type=int,default=2048); ap.add_argument("--accum",type=int,default=8) ap.add_argument("--rank",type=int,default=128); ap.add_argument("--lr",type=float,default=1e-4) ap.add_argument("--steps",type=int,default=0); ap.add_argument("--warmup",type=int,default=20) ap.add_argument("--eval-blocks",type=int,default=96); ap.add_argument("--eval-every",type=int,default=50) ap.add_argument("--save-every",type=int,default=25); ap.add_argument("--resume",action="store_true"); ap.add_argument("--load-4bit",action="store_true") ap.add_argument("--adapter-init",default=None) A=ap.parse_args(); os.makedirs(A.out,exist_ok=True) rank=int(os.environ.get("LOCAL_RANK",0)); world=int(os.environ.get("WORLD_SIZE",1)) ddp=world>1 if ddp: torch.distributed.init_process_group("nccl"); torch.cuda.set_device(rank) dev=rank tok=AutoTokenizer.from_pretrained(A.model) # cache keyed on corpus content+seq, NOT the out dir (the v5 poisoned-cache lesson) _h=hashlib.sha256() with open(A.corpus,'rb') as _f: for _b in iter(lambda:_f.read(1<<22), b''): _h.update(_b) _h.update(str(os.path.getsize(A.corpus)).encode()) ch=_h.hexdigest()[:12] cache=f"/workspace/aditi/corpus/blocks_{ch}_seq{A.seq}.pt" ids,lab=build_blocks(tok,A.corpus,A.seq,cache) g=torch.Generator().manual_seed(1947); perm=torch.randperm(len(ids),generator=g) ne=min(A.eval_blocks,max(8,len(ids)//10)) eidx,tidx=perm[:ne],perm[ne:] steps=A.steps or max(1,len(tidx)//A.accum) if rank==0: print(f"[data] train {len(tidx)} eval {len(eidx)} blocks | " f"one epoch = {len(tidx)//A.accum} steps | running {steps}",flush=True) if A.load_4bit: bnb=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True) m=AutoModelForCausalLM.from_pretrained(A.model,quantization_config=bnb,dtype=torch.bfloat16, device_map={"":dev},attn_implementation="sdpa") if rank==0: print("[base] NF4 4-bit (grid != q4_0 deployment target)",flush=True) else: m=AutoModelForCausalLM.from_pretrained(A.model,dtype=torch.bfloat16, device_map={"":dev},attn_implementation="sdpa") for p_ in m.parameters(): p_.requires_grad_(False) if rank==0: print("[base] bf16 QAT weights - adapter learns NO quantisation error",flush=True) try: from liger_kernel.transformers import _apply_liger_kernel_to_instance _apply_liger_kernel_to_instance(model=m) if rank==0: print("[patch] liger fused CE",flush=True) except Exception as e: print("[patch] liger unavailable",e,flush=True) m.config.use_cache=False if A.adapter_init: from peft import PeftModel mm=PeftModel.from_pretrained(m,A.adapter_init,is_trainable=True) if rank==0: print(f"[peft] WARM-START from {A.adapter_init}",flush=True) else: mm=get_peft_model(m,LoraConfig(r=A.rank,lora_alpha=A.rank*2,lora_dropout=0.05,bias="none", task_type="CAUSAL_LM",target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"])) mm.train() ps=[p for p in mm.parameters() if p.requires_grad] if rank==0: print(f"[peft] trainable {sum(p.numel() for p in ps)/1e6:.1f}M rank={A.rank}",flush=True) opt=torch.optim.AdamW(ps,lr=A.lr,betas=(0.9,0.95),weight_decay=0.0) sched=lambda s: min(1.0,(s+1)/max(A.warmup,1))*0.5*(1+math.cos(math.pi*min(1.0,s/steps))) START=0; st=os.path.join(A.out,"last","state.json") if A.resume and os.path.exists(st): START=json.load(open(st))["step"] from peft import PeftModel sd=torch.load(os.path.join(A.out,"last","optimizer.pt"),map_location="cpu") opt.load_state_dict(sd["opt"]); print(f"[resume] from step {START}",flush=True) def save(step): if rank!=0: return d=os.path.join(A.out,"last"); os.makedirs(d,exist_ok=True) mm.save_pretrained(d); tok.save_pretrained(d) torch.save({"opt":opt.state_dict()},os.path.join(d,"optimizer.pt")) json.dump({"step":step},open(os.path.join(d,"state.json"),"w")) @torch.no_grad() def evaluate(): mm.eval(); tot=0.0;k=0 for i in eidx[:ne]: o=mm(input_ids=ids[i:i+1].to(dev),labels=lab[i:i+1].to(dev)) l=o.loss.item() if l==l: tot+=l;k+=1 mm.train(); return tot/max(k,1) t0=time.time(); pos=START*A.accum; mlog=open(os.path.join(A.out,"metrics.jsonl"),"a") for step in range(START,steps): for gi in range(A.accum): i=tidx[(pos)%len(tidx)]; pos+=1 if ddp and (gi%world)!=rank: continue o=mm(input_ids=ids[i:i+1].to(dev),labels=lab[i:i+1].to(dev)) (o.loss/A.accum*(world if ddp else 1)).backward() if ddp: for p in ps: if p.grad is not None: torch.distributed.all_reduce(p.grad) torch.nn.utils.clip_grad_norm_(ps,1.0) for gp in opt.param_groups: gp["lr"]=A.lr*sched(step) opt.step(); opt.zero_grad(set_to_none=True) if rank==0 and (step%5==0 or step==steps-1): el=time.time()-t0; done=step-START+1 tps=done*A.accum*A.seq/max(el,1e-9) pk=torch.cuda.max_memory_allocated(dev)/2**30 print(f"[step {step+1}/{steps}] loss {o.loss.item():.4f} lr {A.lr*sched(step):.2e} " f"{tps:.0f} tok/s peak {pk:.1f}G elapsed {el/60:.1f}m",flush=True) mlog.write(json.dumps({"step":step+1,"loss":o.loss.item(),"tok_s":round(tps),"peak_gib":round(pk,1)})+"\n"); mlog.flush() if (step+1)%A.eval_every==0 and rank==0: print(f"[eval] step {step+1} held-out loss {evaluate():.4f}",flush=True) if (step+1)%A.save_every==0: save(step+1) save(steps) if rank==0: print(f"[DONE] {steps} steps in {(time.time()-t0)/3600:.2f}h " f"peak {torch.cuda.max_memory_allocated(dev)/2**30:.1f} GiB",flush=True) if __name__=="__main__": main()