File size: 8,673 Bytes
e7d8c8b 6c8db72 e7d8c8b 6c8db72 e7d8c8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | #!/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()
|