#!/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()