"""Gemma-4 26B-A4B (MoE, model_type gemma4) TP=8 prefill+decode via NxD ModelBuilder. Reuses the 31B tp_mb.py recipe wholesale (ModelBuilder single-rank compile, mixed sliding/global attention shard/replicate, ScatterKV device cache, layer_scalar buffers, chat-template prompt, device prefill+decode) and swaps ONLY the FFN for the MoE dual-path. 26B-A4B is a DUAL-PATH FFN per layer: a shared dense `mlp` (intermediate 2112) in PARALLEL with a 128-expert MoE (top-8, moe_intermediate 704), combined `dense+moe` through 4 feed-forward layernorms. HF's decoder-layer forward + `Gemma4TextRouter` are kept UNCHANGED (router replicated); we replace only `self.experts` (whose HF forward is a sparse gather/scatter loop that won't trace) with `DenseExperts`: * all-experts-DENSE compute (all 128 experts on every token, weighted by the top-8 router weight, 0 for non-selected -> exact match to HF's sparse top-8, static-shape + traceable). * expert weights mapped onto TWO standard parallel linears ModelBuilder auto-shards: gate_up: [E*2*I, H] ColumnParallelLinear (rank r gets experts [16r:16r+16]) down: [H, E*I] RowParallelLinear (input-sharded by expert-block -> all-reduce) -> ~5.7GB experts/rank (can't replicate 45.6GB > 16GB/core). Set MODEL_DIR=/data/real-gemma4-26B-A4B-it MB_WDTYPE=bf16 TP_DEGREE=8 KV_MAX=256 KV_BUCKET=64.""" import sys, os, types, time m=types.ModuleType("transformers.utils.fx"); m.HFTracer=object; m.symbolic_trace=None; sys.modules["transformers.utils.fx"]=m import multiprocessing MP=os.environ.get("MODEL_DIR","/data/real-gemma4-26B-A4B-it"); TP=int(os.environ.get("TP_DEGREE","8")) MAX=int(os.environ.get("KV_MAX","256")); BUCKET=int(os.environ.get("KV_BUCKET","64")) CARGS="--model-type transformer --auto-cast all --auto-cast-type bf16" def _kv_rank_width(nkv): return nkv//TP if nkv%TP==0 else nkv def _discover(): import torch from transformers import Gemma4ForConditionalGeneration mm=Gemma4ForConditionalGeneration.from_pretrained(MP,torch_dtype=torch.float32,attn_implementation="eager"); mm.eval() lang=mm.model.language_model; cfg=lang.config NONSHARED,LINFO=[],{} for i,lyr in enumerate(lang.layers[:cfg.num_hidden_layers]): a=lyr.self_attn if not getattr(a,"is_kv_shared_layer",False): hd=a.head_dim; NONSHARED.append(i); LINFO[i]=(a.k_proj.out_features//hd,hd) softcap=getattr(mm.config.text_config,"final_logit_softcapping",None) return mm,lang,mm.lm_head,softcap,NONSHARED,LINFO,cfg.sliding_window def _structure(): import torch from transformers import Gemma4ForConditionalGeneration, AutoConfig cfg=AutoConfig.from_pretrained(MP) cfg._attn_implementation="eager" if hasattr(cfg,"text_config"): cfg.text_config._attn_implementation="eager" mm=Gemma4ForConditionalGeneration(cfg); mm.eval() lang=mm.model.language_model; lc=lang.config NONSHARED,LINFO=[],{} for i,lyr in enumerate(lang.layers[:lc.num_hidden_layers]): a=lyr.self_attn if not getattr(a,"is_kv_shared_layer",False): hd=a.head_dim; NONSHARED.append(i); LINFO[i]=(a.k_proj.out_features//hd,hd) softcap=getattr(mm.config.text_config,"final_logit_softcapping",None) return mm,lang,mm.lm_head,softcap,NONSHARED,LINFO,lc.sliding_window def build_module(): import torch from neuronx_distributed.parallel_layers.layers import ColumnParallelLinear, RowParallelLinear, SPMDRank from neuronx_distributed.parallel_layers.mappings import scatter_to_process_group_spmd mm,lang,lm_head,softcap,NONSHARED,LINFO,SW=_structure() lc=lang.config N_EXP=lc.num_experts; MOE_I=lc.moe_intermediate_size; H=lc.hidden_size; TOPK=lc.top_k_experts class GeluTanh(torch.nn.Module): def forward(s,x): return 0.5*x*(1.0+torch.tanh(0.7978845608028654*(x+0.044715*x*x*x))) for mod in lang.modules(): if hasattr(mod,"act_fn"): mod.act_fn=GeluTanh() WDT=torch.bfloat16 if os.environ.get("MB_WDTYPE","bf16")=="bf16" else torch.float32 def col(o): return ColumnParallelLinear(o.in_features,o.out_features,bias=False,gather_output=False,dtype=WDT) def row(o): return RowParallelLinear(o.in_features,o.out_features,bias=False,input_is_parallel=True,dtype=WDT) class DenseExperts(torch.nn.Module): """All-experts-dense MoE, expert-parallel via two standard parallel linears. forward(x, top_k_index, top_k_weights) -> [T, H] (matches HF Gemma4TextExperts signature). top_k_weights already carries the router's renorm + per_expert_scale.""" def __init__(s): super().__init__() s.gate_up=ColumnParallelLinear(H, N_EXP*2*MOE_I, bias=False, gather_output=False, dtype=WDT) s.down=RowParallelLinear(N_EXP*MOE_I, H, bias=False, input_is_parallel=True, dtype=WDT) s.act=GeluTanh(); s.epr=N_EXP//TP s.spmd_rank=SPMDRank(TP) # runtime rank tensor for the SPMD scatter (checkpoint = arange(TP)) def forward(s, x, top_k_index, top_k_weights): T=x.shape[0] ar=torch.arange(N_EXP, device=x.device) onehot=(top_k_index.unsqueeze(-1)==ar).to(x.dtype) # [T,K,E] Wd=(top_k_weights.unsqueeze(-1).to(x.dtype)*onehot).sum(dim=1) # [T,E] dense, 0 for non-top-k # SPMD-aware scatter: standalone scatter_to_tensor_model_parallel_region bakes rank 0's slice # into the single-rank trace -> every rank weights experts 0..15 while gate_up computes its # OWN experts (misaligned garbage). This uses the runtime rank tensor to pick Wd[:,16r:16r+16]. Wl=scatter_to_process_group_spmd(Wd, 1, s.spmd_rank.get_rank()) # [T, epr] rank-local experts gu=s.gate_up(x).view(T, s.epr, 2*MOE_I) # [T,epr,2I] gate,up=gu.chunk(2, dim=-1) # [T,epr,I] each h=(s.act(gate)*up)*Wl.unsqueeze(-1) # weight before down h=h.reshape(T, s.epr*MOE_I) # [T, epr*I] return s.down(h) # RowParallel all-reduce -> [T,H] nshard=nrepl=nmoe=0 for lyr in lang.layers[:lc.num_hidden_layers]: a=lyr.self_attn; hd=a.head_dim k_out=a.k_proj.out_features if getattr(a,"k_proj",None) is not None else None nkv=(k_out//hd) if k_out is not None else None if nkv is not None and nkv%TP==0: a.q_proj=col(a.q_proj); a.o_proj=row(a.o_proj); a.k_proj=col(a.k_proj) if getattr(a,"v_proj",None) is not None: a.v_proj=col(a.v_proj) nshard+=1 else: nrepl+=1 mp=lyr.mlp; mp.gate_proj=col(mp.gate_proj); mp.up_proj=col(mp.up_proj); mp.down_proj=row(mp.down_proj) # shared dense MLP if getattr(lyr,"enable_moe_block",False) and hasattr(lyr,"experts"): lyr.experts=DenseExperts(); nmoe+=1 # router stays intact (replicated) print(f"build_module: {nshard} sharded-attn, {nrepl} replicated-attn, {nmoe} MoE layers, TP={TP} (E={N_EXP} K={TOPK} I={MOE_I})",flush=True) import glob from safetensors import safe_open lsv={} for fp in sorted(glob.glob(MP+"/*.safetensors")): with safe_open(fp,framework="pt") as sf: for k in sf.keys(): if k.startswith("model.language_model.layers.") and k.endswith(".layer_scalar"): lsv[int(k.split("layers.")[1].split(".")[0])]=sf.get_tensor(k) with torch.no_grad(): for i,lyr in enumerate(lang.layers[:lc.num_hidden_layers]): if hasattr(lyr,"layer_scalar") and i in lsv: lyr.layer_scalar.copy_(lsv[i].to(lyr.layer_scalar.dtype)) print(f" loaded {len(lsv)} layer_scalar buffers",flush=True) lang.embed_tokens=torch.nn.Embedding(2, H) NK=len(NONSHARED); W={i:_kv_rank_width(LINFO[i][0]) for i in NONSHARED} class ScatterKV: is_compileable=False def __init__(s,kb,vb,oh): s.key={i:kb[j] for j,i in enumerate(NONSHARED)}; s.val={i:vb[j] for j,i in enumerate(NONSHARED)}; s.oh=oh def _wr(s,buf,x): oh=s.oh[0,0]; scat=torch.einsum('bhsd,ms->bhmd',x,oh); r=oh.sum(-1).view(1,1,MAX,1); return buf*(1.0-r)+scat def update(s,k,v,idx,*a,**kw): s.key[idx]=s._wr(s.key[idx],k); s.val[idx]=s._wr(s.val[idx],v); return s.key[idx],s.val[idx] def get_seq_length(s,*a,**k): return 0 def export(s): return [s.key[i] for i in NONSHARED],[s.val[i] for i in NONSHARED] class Wrap(torch.nn.Module): def __init__(s): super().__init__(); s.lang=lang; s.head=lm_head s.kbuf=torch.nn.ParameterList([torch.nn.Parameter(torch.zeros(1,W[i],MAX,LINFO[i][1]),requires_grad=False) for i in NONSHARED]) s.vbuf=torch.nn.ParameterList([torch.nn.Parameter(torch.zeros(1,W[i],MAX,LINFO[i][1]),requires_grad=False) for i in NONSHARED]) def forward(s, ie, position_ids, onehot, full_mask, slide_mask): cache=ScatterKV(list(s.kbuf),list(s.vbuf),onehot) out=s.lang(inputs_embeds=ie,position_ids=position_ids, attention_mask={"full_attention":full_mask,"sliding_attention":slide_mask},use_cache=True,past_key_values=cache) lg=s.head(out.last_hidden_state); ks,vs=cache.export(); return (lg,)+tuple(ks)+tuple(vs) w=Wrap().eval() aliases={} for j in range(NK): aliases[w.kbuf[j]]=1+j for j in range(NK): aliases[w.vbuf[j]]=1+NK+j return w, aliases def checkpoint_loader(): """Full fp32 state dict keyed to the Wrap. Reshape the fused 3D expert weights into the 2D ColumnParallel(gate_up)/RowParallel(down) layouts DenseExperts expects.""" import torch, glob, re from safetensors.torch import load_file raw={} for fp in sorted(glob.glob(MP+"/*.safetensors")): raw.update(load_file(fp)) sd={} exp_re=re.compile(r"layers\.(\d+)\.experts\.(gate_up_proj|down_proj)$") for k,v in raw.items(): if k=="model.language_model.embed_tokens.weight": sd["head.weight"]=v.to(torch.float32) elif k.startswith("model.language_model."): kk=k[len("model.language_model."):] mo=exp_re.match(kk) if mo: i=mo.group(1) if mo.group(2)=="gate_up_proj": # [E, 2I, H] -> [E*2I, H] E,TWOI,Hh=v.shape sd[f"lang.layers.{i}.experts.gate_up.weight"]=v.reshape(E*TWOI, Hh).to(torch.float32) else: # down_proj [E, H, I] -> [H, E*I] E,Hh,I=v.shape sd[f"lang.layers.{i}.experts.down.weight"]=v.permute(1,0,2).reshape(Hh, E*I).to(torch.float32) else: sd["lang."+kk]=v.to(torch.float32) elif k.startswith("lm_head."): sd["head."+k[len("lm_head."):]]=v.to(torch.float32) # SPMDRank checkpoint: arange(TP) sharded dim0 -> each rank gets its own rank number at runtime. moe_layers={int(re.match(r"lang\.layers\.(\d+)\.experts\.gate_up\.weight",kk).group(1)) for kk in sd if re.match(r"lang\.layers\.(\d+)\.experts\.gate_up\.weight",kk)} for i in sorted(moe_layers): sd[f"lang.layers.{i}.experts.spmd_rank.rank"]=torch.arange(TP, dtype=torch.int32) return sd def _inputs(lang, ids, SW, NEG, positions): import torch ids_t=torch.tensor([ids]) with torch.no_grad(): ie=lang.embed_tokens(ids_t) seq=len(positions); ar=torch.arange(MAX); pos=torch.tensor([positions],dtype=torch.long) oh=(ar.view(MAX,1)==pos.view(1,seq)).view(1,1,MAX,seq).to(torch.float32) q=pos.view(seq,1); mm2=ar.view(1,MAX) full=torch.where(mm2<=q,0.0,NEG).view(1,1,seq,MAX) slide=torch.where((mm2<=q)&(mm2>q-SW),0.0,NEG).view(1,1,seq,MAX) return ie,pos,oh,full,slide def main(): import torch from transformers import AutoTokenizer, DynamicCache from neuronx_distributed.trace.model_builder import ModelBuilder, BaseModelInstance NEG=torch.finfo(torch.float32).min tok=AutoTokenizer.from_pretrained(MP) _mm,rlang,_h,softcap,NONSHARED,LINFO,SW=_discover() print(f"26B-A4B discover: {len(NONSHARED)} non-shared layers, head_dims={sorted({LINFO[i][1] for i in NONSHARED})}, softcap={softcap}",flush=True) if not getattr(tok,"chat_template",None): try: import subprocess,json from huggingface_hub import hf_hub_download raw=subprocess.check_output(["aws","secretsmanager","get-secret-value","--region","us-east-1","--secret-id","hf_token","--query","SecretString","--output","text"]).decode().strip() try: tv=json.loads(raw); tv=tv.get("hf_token") or list(tv.values())[0] except Exception: tv=raw tok.chat_template=open(hf_hub_download(os.environ.get("HF_REPO","google/gemma-4-26B-A4B-it"),"chat_template.jinja",token=tv)).read() except Exception as e: print("chat_template fetch failed:",e,flush=True) if getattr(tok,"chat_template",None): d=tok.apply_chat_template([{"role":"user","content":"What is the capital of France?"}],add_generation_prompt=True) prompt=d if isinstance(d,list) else d["input_ids"] else: prompt=[2,105,2364,107,3689,563,506,5279,529,7001,236881,106,107,105,4368,107,100,45518,107,101] n0=len(prompt); assert n0<=BUCKET pre_in=_inputs(rlang, prompt+[0]*(BUCKET-n0), SW, NEG, list(range(BUCKET))) dec_in=_inputs(rlang, [prompt[-1]], SW, NEG, [n0]) load_from=os.environ.get("MB_LOAD") if load_from: print("loading saved model from",load_from,flush=True) model=torch.jit.load(load_from) model.nxd_model.initialize_with_saved_weights(torch.tensor([0],dtype=torch.int32)); print("MB_LOADED",flush=True) else: class GemmaInstance(BaseModelInstance): def __init__(self): self.module=None; self.input_output_aliases=[{}] def load_module(self): self.module, aliases = build_module() self.input_output_aliases=[aliases] def get(self, bucket_rank, **kwargs): return self.module, self.input_output_aliases[0] inst=GemmaInstance() print("building ModelBuilder ...",flush=True) mb=ModelBuilder(router=None, tp_degree=TP, checkpoint_loader=checkpoint_loader, compiler_workdir="/data/mb_wd") mb.add("prefill", inst, [pre_in], compiler_args=CARGS) mb.add("decode", inst, [dec_in], compiler_args=CARGS) print("tracing (ModelBuilder, single-rank compile) ...",flush=True) model=mb.trace(initialize_model_weights=True) print("MB_TRACED",flush=True) save_to=os.environ.get("MB_SAVE") if save_to: print("saving model to",save_to,flush=True); torch.jit.save(model, save_to); print("MB_SAVED",save_to,flush=True) if os.environ.get("SKIP_VALIDATE")=="1": print("SKIP_VALIDATE=1 -> compile-only run complete",flush=True); print("TPMB_OK",flush=True); return head=_mm.lm_head ec=_mm.generation_config.eos_token_id; EOS=set(ec) if isinstance(ec,(list,tuple)) else {ec} class _G(torch.nn.Module): def forward(s,x): return 0.5*x*(1.0+torch.tanh(0.7978845608028654*(x+0.044715*x*x*x))) for mod in rlang.modules(): if hasattr(mod,"act_fn"): mod.act_fn=_G() def _sc2(lg): return softcap*torch.tanh(lg/softcap) if softcap else lg def embR(ids): idt=torch.tensor([ids]) with torch.no_grad(): ie=rlang.embed_tokens(idt) return ie def cpuref(): ie=embR(prompt); c=DynamicCache() with torch.no_grad(): out=rlang(inputs_embeds=ie,attention_mask=torch.ones(1,n0,dtype=torch.long),use_cache=True,past_key_values=c) first=int(_sc2(head(out.last_hidden_state))[0,n0-1].argmax()) KS=[c.layers[i].keys for i in NONSHARED]; VS=[c.layers[i].values for i in NONSHARED] class SKV: is_compileable=False def __init__(s,kb,vb,oh): s.key={i:kb[j] for j,i in enumerate(NONSHARED)}; s.val={i:vb[j] for j,i in enumerate(NONSHARED)}; s.oh=oh def update(s,k,v,idx,*a,**kw): s.key[idx]=s.key[idx]*(1.0-s.oh)+k*s.oh; s.val[idx]=s.val[idx]*(1.0-s.oh)+v*s.oh; return s.key[idx],s.val[idx] def get_seq_length(s,*a,**k): return 0 def hp(pos): ar=torch.arange(MAX); oh=(ar==pos).view(1,1,MAX,1).float(); valid=ar<=pos fm=torch.where(valid,0.0,NEG).view(1,1,1,MAX); sm=torch.where(valid&(ar>pos-SW),0.0,NEG).view(1,1,1,MAX) return torch.tensor([[pos]]),oh,fm,sm kb=[torch.zeros(1,LINFO[i][0],MAX,LINFO[i][1]) for i in NONSHARED]; vb=[torch.zeros(1,LINFO[i][0],MAX,LINFO[i][1]) for i in NONSHARED] for j in range(len(NONSHARED)): kb[j][:,:,:n0,:]=KS[j][:,:,:n0,:]; vb[j][:,:,:n0,:]=VS[j][:,:,:n0,:] seq=[first]; cur=n0 for _ in range(30): if seq[-1] in EOS: break e1=embR([seq[-1]]); pi,oh,fm,sm=hp(cur); cache=SKV(kb,vb,oh) with torch.no_grad(): out=rlang(inputs_embeds=e1,position_ids=pi,attention_mask={"full_attention":fm,"sliding_attention":sm},use_cache=True,past_key_values=cache) lg=_sc2(head(out.last_hidden_state)) kb=[cache.key[i] for i in NONSHARED]; vb=[cache.val[i] for i in NONSHARED] seq.append(int(lg[0,0].argmax())); cur+=1 return seq DEVICE_ONLY=os.environ.get("DEVICE_ONLY")=="1" if DEVICE_ONLY: cpu_seq=None; print("DEVICE_ONLY=1 -> skipping CPU reference",flush=True) else: print("=== CPU ref ===",flush=True); cpu_seq=cpuref() def dcall(a): r=model(*a); return r[0] if isinstance(r,(tuple,list)) else r pad=prompt+[0]*(BUCKET-n0) print("=== DEVICE (bucketed) ===",flush=True) dcall(_inputs(rlang,pad,SW,NEG,list(range(BUCKET)))) t0=time.time(); lg=dcall(_inputs(rlang,pad,SW,NEG,list(range(BUCKET)))); pf=time.time()-t0 first=int(lg[0,n0-1].argmax()); seq=[first]; cur=n0 for _ in range(30): if seq[-1] in EOS: break l1=dcall(_inputs(rlang,[seq[-1]],SW,NEG,[cur])); seq.append(int(l1[0,0].argmax())); cur+=1 dev_txt=tok.decode([x for x in seq if x not in EOS],skip_special_tokens=True) if cpu_seq is not None: print("CPU GEN:",repr(tok.decode([x for x in cpu_seq if x not in EOS],skip_special_tokens=True)),flush=True) print("SEQ_MATCH", cpu_seq==seq, flush=True) else: print("DEVICE_PARIS", "paris" in dev_txt.lower(), flush=True) print("DEV GEN:",repr(dev_txt),flush=True) print(f"DEVICE PREFILL first-token: {pf*1000:.0f} ms",flush=True) print("TPMB_OK",flush=True) if __name__=="__main__": if os.environ.get("_TP_CHILD")!="1": os.environ["_TP_CHILD"]="1"; multiprocessing.set_start_method("spawn",force=True); main()