File size: 5,653 Bytes
6c8db72 | 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 | #!/usr/bin/env python
"""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)
# compressed_tensors: scale = amax / (bit_range/2); int4 range is [-8,7]
# so bit_range = 15 and the divisor is 7.5, NOT 7.0. Verified bit-exact against
# compressed_tensors.fake_quantize on real layers (rel-err 0.000000).
# Was 7.0 until 2026-08-21 -> trained against a grid 8.6% off the deployed one.
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) # <-- the key step
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()
|