| |
| """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) |
| 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) |
| |
| _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() |
|
|
