| |
| """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) |
| |
| 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() |
|
|
