FunctionGemma export materializes the tied embedding twice (77% of the weight archive)

[michalharakal]

Summary

The FunctionGemma-270M compiled export emits the tied vocab embedding as two independent weight globals, roughly doubling the archive and every in-memory copy during export.

Evidence

Parsed from /tmp/fgemma-export/gemma-gen.mlir (produced by FunctionGemmaExport.export):

globals: 324   total elements: 436,111,680   = 1.62 GiB FP32   (832 MiB bf16)

top globals by size:
  t2044    262153x640    167,777,920   (38.5%)
  t0       262153x640    167,777,920   (38.5%)
  t98        2048x640      1,310,720   ( 0.3%)
  ...

t0 (input token embedding) and t2044 (output logit projection) are the same tied weight, emitted twice. Together they are 77% of the archive. The 832 MiB figure matches the on-disk gemma.safetensors exactly (872,248,326 bytes), so this is measured, not estimated.

The model itself is ~268.3M params; the archive carries 436.1M elements. The 167.8M difference is exactly the duplicate.

Impact

• Archive size: 832 MiB bf16 where ~512 MiB would do. On the 1.9 GB SL2610 board this is a real constraint — it is the reason the int8 path (415 MiB) was needed.
• Export memory: the duplicate is carried in every one of the ~3 near-simultaneous FP32 copies during export, adding ~640 MiB per copy. It is a significant contributor to the ~4.3 GiB peak live heap that forced the test-heap bump in kgemma: heavy full-model-trace tests OOM at the 4 GB default heap #250 / PR fix(kgemma): raise heavy-trace test heap to 12g and make skips honest #259.

Suggested fix

Emit the tied embedding once and have both consumers reference the same util.global / external parameter. Gemma ties token_embd to the output head by construction, so the second materialization carries no information.

Worth confirming first whether the duplication originates in the DSL graph (two distinct nodes holding the same tensor) or in the StableHLO externalization step assigning two parameter keys to one buffer — the fix differs.

Related

• FunctionGemma export: convert MLIR-text rewrites to StableHLO/DAG transforms #248 — FunctionGemma export post-processing (same file, FunctionGemmaExport.kt)
• KV-cache trace fidelity: shared-cache variants still bake K=V=0 under embedConstants #194 — embedConstants trace fidelity
• kgemma: heavy full-model-trace tests OOM at the 4 GB default heap #250 / PR fix(kgemma): raise heavy-trace test heap to 12g and make skips honest #259 — the test-heap bump this waste contributed to
• Shared narrow-float layer: complete BF16, add FP16 on the same SPI SKaiNET#884 — upstream narrow-float layer; halving the copy width is complementary to removing this duplicate

🤖 Generated with Claude Code