| |
| """ |
| Aether Phase-1 CACHED alignment trainer — trains from pre-extracted features with |
| NO ENCODERS LOADED. This is the portable path: base LM (or 4-bit QLoRA of it) + the |
| tiny projectors fit a 16GB card, so alignment/adapters run on 6900XT / Kaggle-T4 / Colab. |
| |
| Consumes /work/slivers_cache/{index.jsonl, <modality>/<id>.pt} from preextract.py. |
| Proves loss drops on REAL cached features across ALL THREE modalities at once. |
| """ |
| import os, sys, json, torch, torch.nn as nn |
| sys.path.insert(0,"/work") |
| import forward as F |
| from transformers import AutoModelForImageTextToText, AutoTokenizer |
| dev=F.dev; TXT=F.TXT; CACHE="/work/slivers_cache" |
|
|
| class CachedAether(nn.Module): |
| """Base LM + projectors + expanded vocab. NO encoders (features come pre-extracted).""" |
| def __init__(self): |
| super().__init__() |
| self.base=AutoModelForImageTextToText.from_pretrained("/work/base",torch_dtype=torch.bfloat16,trust_remote_code=True) |
| self.tok=AutoTokenizer.from_pretrained("/work/base",trust_remote_code=True) |
| old=len(self.tok); self.tok.add_special_tokens({"additional_special_tokens":F.build_special()}) |
| self.new_vocab=len(self.tok); self.base.resize_token_embeddings(self.new_vocab); self.new_row_lo=old |
| self.visual_proj=F.Projector(3200,TXT); self.audio_proj=F.Projector(1280,TXT) |
| self.mv_proj=F.Projector(3200,TXT); self.geom_proj=F.Projector(8,TXT) |
| self.cam_pose=nn.Parameter(torch.zeros(1,1,TXT)) |
| self.id={k:self.tok.convert_tokens_to_ids(k) for k in ["<img>","<audio>","<3d_app>","<3d_geom>"]} |
| def trainable(self): |
| for m in (self.visual_proj,self.audio_proj,self.mv_proj,self.geom_proj): yield from m.parameters() |
| yield self.cam_pose |
|
|
| def apply_freeze(m): |
| for p in m.parameters(): p.requires_grad_(False) |
| for p in m.trainable(): p.requires_grad_(True) |
| lo=m.new_row_lo |
| for w in (m.base.get_input_embeddings().weight, m.base.get_output_embeddings().weight): |
| w.requires_grad_(True); w.register_hook(lambda g,lo=lo:(g.__setitem__(slice(0,lo),0) or g)) |
|
|
| PROJ={"vision":("visual_proj","<img>"),"audio":("audio_proj","<audio>"),"geom":("geom_proj","<3d_geom>")} |
|
|
| def splice(m, ids, idxpos, vproj): |
| emb=m.base.get_input_embeddings()(ids.to(dev)) |
| repl=torch.zeros_like(emb); repl[0,idxpos]=vproj.to(emb.dtype) |
| mask=torch.zeros(emb.shape[:2],dtype=torch.bool,device=dev); mask[0,idxpos]=True |
| return torch.where(mask.unsqueeze(-1),repl,emb) |
|
|
| if __name__=="__main__": |
| torch.manual_seed(0) |
| m=CachedAether().to(dev) |
| for mod in (m.visual_proj,m.audio_proj,m.mv_proj,m.geom_proj): mod.to(torch.bfloat16) |
| m.cam_pose.data=m.cam_pose.data.to(torch.bfloat16) |
| apply_freeze(m) |
| print("[cached] base+projectors loaded, NO ENCODERS. VRAM GB:",round(torch.cuda.memory_allocated()/1e9,1)) |
| items=json.load(open(f"{CACHE}/index.jsonl")) |
| print(f"[cached] {len(items)} pre-extracted items:",{r['modality'] for r in items}) |
| T=m.tok |
| def make(r): |
| projname,ptok=PROJ[r["modality"]]; pid=m.id[ptok]; proj=getattr(m,projname) |
| blob=torch.load(r["feat"],map_location=dev) |
| if r["modality"]=="geom": |
| feat=blob["feat"].to(dev); xyz=blob["xyz"].to(dev); v=proj(feat.to(torch.bfloat16))+m.cam_pose.squeeze(0) |
| else: |
| feat=blob.to(dev); xyz=None; v=proj(feat.to(torch.bfloat16)) |
| cid=T(r["caption"],add_special_tokens=False).input_ids |
| ids=torch.tensor([pid]*v.shape[0]+cid)[None] |
| pos_ids=(ids[0]==pid).nonzero(as_tuple=True)[0] |
| labels=torch.tensor([-100]*v.shape[0]+cid)[None].to(dev) |
| ttype=("3d_geom" if r["modality"]=="geom" else ("audio" if r["modality"]=="audio" else "img")) |
| types=[ttype]*v.shape[0]+["text"]*len(cid) |
| pos=F.build_mrope(types,geom_xyz=xyz).to(dev) |
| return ids,pos_ids,v,labels,pos |
| opt=torch.optim.AdamW([p for p in m.parameters() if p.requires_grad],lr=1e-3) |
| STEPS=int(os.environ.get("STEPS","200")) |
| print(f"step loss (STEPS={STEPS})") |
| for step in range(STEPS): |
| opt.zero_grad(); tot=0.0 |
| for r in items: |
| ids,pos_ids,v,labels,pos=make(r) |
| out=m.base(inputs_embeds=splice(m,ids,pos_ids,v),position_ids=pos,labels=labels,use_cache=False) |
| out.loss.backward(); tot+=out.loss.item() |
| opt.step() |
| if step%20==0 or step==STEPS-1: print(f"{step:4d} {tot/len(items):.4f}") |
| ok = tot/len(items) < 1.0 |
| print("CACHED MULTIMODAL ALIGNMENT OK (no encoders) — loss collapsed" if ok else "!! loss did not collapse") |
|
|
