File size: 7,284 Bytes
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 | #!/usr/bin/env python
"""Tri-modal trainer for Gemma 4 12B Unified on the Google QAT foundation.
Trains:
- LoRA on the shared 48-layer trunk (processes text AND vision AND audio)
- vision_embedder (49.9M) + embed_audio (2.5M) FULLY trainable - they are small and
are the only modality-specific weights that exist in this architecture.
No packing: media makes sequences variable, so micro-batch 1 + accumulation.
Assistant-only masking by incremental re-render (never regex).
"""
import os,sys,json,math,time,random,argparse
import torch
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
from transformers import AutoProcessor, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
def load_rows(paths):
R=[]
for p in paths:
for l in open(p,encoding="utf-8"):
l=l.strip()
if l: R.append(json.loads(l))
return R
def build_sample(proc, rec, max_len):
"""Returns input_ids/labels/media tensors with only assistant tokens supervised."""
msgs=rec["messages"]; media=rec.get("media") or {}
img=None; aud=None
if media.get("image"):
from PIL import Image
try: img=Image.open(media["image"]).convert("RGB")
except Exception: return None
if media.get("audio"):
import soundfile as sf
try: aud,_=sf.read(media["audio"],dtype="float32")
except Exception: return None
user=[c for c in msgs if c["role"]=="user"][0]["content"]
content=[]
if img is not None: content.append({"type":"image"})
if aud is not None: content.append({"type":"audio"})
content.append({"type":"text","text":user})
conv=[{"role":"system","content":[{"type":"text","text":msgs[0]["content"]}]},
{"role":"user","content":content}]
asst=[c for c in msgs if c["role"]=="assistant"][0]["content"]
kw={}
if img is not None: kw["images"]=[img]
if aud is not None: kw["audio"]=[aud]
try:
prefix=proc.apply_chat_template(conv, add_generation_prompt=True, tokenize=False)
full=prefix+asst+(proc.tokenizer.eos_token or "")
e_full=proc(text=full, return_tensors="pt", **kw)
n_pre=len(proc(text=prefix, return_tensors="pt", **kw)["input_ids"][0])
except Exception:
return None
ids=e_full["input_ids"][0]
if len(ids)>max_len or len(ids)<=n_pre: return None
lab=ids.clone(); lab[:n_pre]=-100
out={"input_ids":ids.unsqueeze(0),"labels":lab.unsqueeze(0),
"attention_mask":e_full.get("attention_mask",torch.ones_like(ids).unsqueeze(0))}
for k in ("pixel_values","image_position_ids","mm_token_type_ids",
"input_features","audio_position_ids","input_features_mask"):
if k in e_full: out[k]=e_full[k]
return out
def main():
ap=argparse.ArgumentParser()
ap.add_argument("--corpus",nargs="+",required=True); ap.add_argument("--out",required=True)
ap.add_argument("--model",default="/workspace/models/gemma4-qat-unquantized")
ap.add_argument("--rank",type=int,default=32); ap.add_argument("--lr",type=float,default=1e-4)
ap.add_argument("--mm-lr",type=float,default=2e-5)
ap.add_argument("--accum",type=int,default=8); ap.add_argument("--steps",type=int,default=400)
ap.add_argument("--warmup",type=int,default=20); ap.add_argument("--max-len",type=int,default=2048)
ap.add_argument("--save-every",type=int,default=50); ap.add_argument("--adapter-init",default=None)
A=ap.parse_args(); os.makedirs(A.out,exist_ok=True)
proc=AutoProcessor.from_pretrained(A.model)
m=AutoModelForCausalLM.from_pretrained(A.model,dtype=torch.bfloat16,
device_map={"":0},attn_implementation="sdpa")
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 as e: print("[patch] liger off:",str(e)[:80],flush=True)
m.config.use_cache=False
for p in m.parameters(): p.requires_grad_(False)
lc=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"])
if A.adapter_init:
from peft import PeftModel
mm=PeftModel.from_pretrained(m,A.adapter_init,is_trainable=True)
print(f"[peft] warm-started from {A.adapter_init}",flush=True)
else:
mm=get_peft_model(m,lc)
# modality weights: FULLY trainable (only 52M params, and the sole modality-specific tensors)
mmp=[]
for n,p in mm.named_parameters():
if any(k in n for k in ("vision_embedder","embed_vision","embed_audio")):
p.requires_grad_(True); mmp.append(p)
lora=[p for n,p in mm.named_parameters() if p.requires_grad and "lora_" in n]
print(f"[peft] lora {sum(p.numel() for p in lora)/1e6:.1f}M "
f"modality {sum(p.numel() for p in mmp)/1e6:.1f}M trainable",flush=True)
mm.train()
opt=torch.optim.AdamW([{"params":lora,"lr":A.lr},{"params":mmp,"lr":A.mm_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)))
rows=load_rows(A.corpus); random.Random(1947).shuffle(rows)
print(f"[data] {len(rows)} multimodal+text records",flush=True)
t0=time.time(); i=0; used=0; skipped=0
for step in range(A.steps):
for _ in range(A.accum):
s=None
while s is None:
s=build_sample(proc,rows[i%len(rows)],A.max_len); i+=1
if s is None: skipped+=1
if i>len(rows)*3: break
if s is None: continue
s={k:(v.to(0) if hasattr(v,"to") else v) for k,v in s.items()}
out=mm(**s); (out.loss/A.accum).backward(); used+=1
torch.nn.utils.clip_grad_norm_(lora+mmp,1.0)
for gi,base in zip(opt.param_groups,(A.lr,A.mm_lr)): gi["lr"]=base*sched(step)
opt.step(); opt.zero_grad(set_to_none=True)
if step%5==0 or step==A.steps-1:
pk=torch.cuda.max_memory_allocated(0)/2**30
print(f"[step {step+1}/{A.steps}] loss {out.loss.item():.4f} "
f"used={used} skipped={skipped} peak {pk:.1f}G {(time.time()-t0)/60:.1f}m",flush=True)
if (step+1)%A.save_every==0:
d=os.path.join(A.out,"last"); os.makedirs(d,exist_ok=True)
mm.save_pretrained(d); proc.save_pretrained(d)
torch.save({n:p.detach().cpu() for n,p in mm.named_parameters()
if any(k in n for k in ("vision_embedder","embed_vision","embed_audio"))},
os.path.join(d,"modality_weights.pt"))
json.dump({"step":step+1},open(os.path.join(d,"state.json"),"w"))
d=os.path.join(A.out,"last"); os.makedirs(d,exist_ok=True)
mm.save_pretrained(d); proc.save_pretrained(d)
torch.save({n:p.detach().cpu() for n,p in mm.named_parameters()
if any(k in n for k in ("vision_embedder","embed_vision","embed_audio"))},
os.path.join(d,"modality_weights.pt"))
print(f"[DONE] {A.steps} steps in {(time.time()-t0)/3600:.2f}h used={used} skipped={skipped}",flush=True)
if __name__=="__main__": main()
|