| |
| """QAT-AWARE LoRA: fake-quantise the base to the EXACT deployment grid (int4/group-32/ |
| symmetric) before training, so the adapter learns to compensate for int4 rounding rather |
| than being crushed onto the grid afterwards. This is the principled fix for the |
| schema_lock collapse seen at export.""" |
| import os,sys,json,math,time,argparse,hashlib |
| import torch, torch.nn as nn |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from peft import LoraConfig, get_peft_model, PeftModel |
|
|
| @torch.no_grad() |
| def fake_quant_int4_g32(W, group=32): |
| """quantise->dequantise on the deployment grid. Symmetric int4, per-group scales.""" |
| o,i = W.shape |
| if i % group: return W |
| w = W.float().reshape(o, i//group, group) |
| |
| |
| |
| |
| s = w.abs().amax(dim=-1, keepdim=True) / 7.5 |
| s = torch.where(s == 0, torch.ones_like(s), s) |
| q = torch.clamp(torch.round(w / s), -8, 7) |
| return (q * s).reshape(o, i).to(W.dtype) |
|
|
| def apply_fake_quant(model, verbose=True): |
| n=0 |
| for name,mod in model.named_modules(): |
| if not isinstance(mod, nn.Linear): continue |
| if any(k in name for k in ("embed_vision","vision_embedder","embed_audio","lm_head")): continue |
| if "language_model.layers" not in name: continue |
| mod.weight.data = fake_quant_int4_g32(mod.weight.data); n+=1 |
| if verbose: print(f"[qat] fake-quantised {n} Linear layers to int4/group-32", flush=True) |
| return n |
|
|
| 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("--adapter-init",default=None) |
| ap.add_argument("--rank",type=int,default=32); ap.add_argument("--lr",type=float,default=8e-5) |
| ap.add_argument("--seq",type=int,default=2048); ap.add_argument("--accum",type=int,default=8) |
| ap.add_argument("--steps",type=int,default=600); ap.add_argument("--warmup",type=int,default=30) |
| ap.add_argument("--save-every",type=int,default=200); ap.add_argument("--eval-every",type=int,default=200) |
| ap.add_argument("--eval-blocks",type=int,default=96) |
| A=ap.parse_args(); os.makedirs(A.out,exist_ok=True) |
| sys.path.insert(0,"/workspace/aditi") |
| from pack_and_train import build_blocks |
| 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()) |
| ids,lab=build_blocks(tok,A.corpus,A.seq,f"/workspace/aditi/corpus/blocks_{h.hexdigest()[:12]}_seq{A.seq}.pt") |
| 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:] |
| print(f"[data] train {len(tidx)} eval {len(eidx)} blocks | running {A.steps} steps",flush=True) |
| m=AutoModelForCausalLM.from_pretrained(A.model,dtype=torch.bfloat16,device_map={"":0}, |
| attn_implementation="sdpa") |
| apply_fake_quant(m) |
| try: |
| from liger_kernel.transformers import _apply_liger_kernel_to_instance |
| _apply_liger_kernel_to_instance(model=m); print("[patch] liger fused CE",flush=True) |
| except Exception: pass |
| m.config.use_cache=False |
| for p in m.parameters(): p.requires_grad_(False) |
| if A.adapter_init: |
| mm=PeftModel.from_pretrained(m,A.adapter_init,is_trainable=True) |
| print(f"[peft] WARM-START {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] |
| print(f"[peft] trainable {sum(p.numel() for p in ps)/1e6:.1f}M",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/A.steps))) |
| def save(step): |
| d=os.path.join(A.out,"last"); os.makedirs(d,exist_ok=True) |
| mm.save_pretrained(d); tok.save_pretrained(d) |
| json.dump({"step":step},open(os.path.join(d,"state.json"),"w")) |
| t0=time.time(); pos=0 |
| for step in range(A.steps): |
| for _ in range(A.accum): |
| i=tidx[pos%len(tidx)]; pos+=1 |
| out=mm(input_ids=ids[i:i+1].to(0),labels=lab[i:i+1].to(0)) |
| (out.loss/A.accum).backward() |
| 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 step%5==0 or step==A.steps-1: |
| el=time.time()-t0 |
| print(f"[step {step+1}/{A.steps}] loss {out.loss.item():.4f} lr {A.lr*sched(step):.2e} " |
| f"{(step+1)*A.accum*A.seq/max(el,1e-9):.0f} tok/s peak " |
| f"{torch.cuda.max_memory_allocated(0)/2**30:.1f}G elapsed {el/60:.1f}m",flush=True) |
| if (step+1)%A.save_every==0: save(step+1) |
| save(A.steps) |
| print(f"[DONE] {A.steps} steps in {(time.time()-t0)/3600:.2f}h",flush=True) |
| if __name__=="__main__": main() |
|
|
