Trace/export pipeline cannot emit multi-billion-param models: Void ops materialize real zeros and constant extraction duplicates all weights

[michalharakal]

Found while building the Gemma 3n StableHLO export in SKaiNET-transformers (the #377 DSL lane / mobile path). The DSL → tape → toComputeGraph → StableHLO pipeline works (proven at 135M–270M by SmolLM2/FunctionGemma) but cannot emit a 4.5B-parameter model on a 48 GB host:

1. VoidTensorOps allocates real dense zero buffers for every recorded op (ShapeOnlyDataFactory.zeros → DenseTensorDataFactory.createFloatTensorData). Worst case: matmulWeightTransposed calls transpose(weight) per projection, allocating a weight-sized zero tensor that the tape then retains — ~9 GB of retained zeros across a 30-layer trace of Gemma 3n E2B (incl. a 2.1 GB zero for the 262k-vocab lm_head transpose). Shape-only propagation shouldn't materialize data.

2. TraceToGraphBuilder.extractFloatArray copies every frozen weight into the graph while the originals stay referenced by the module tree — dense originals (~10 GB for the E2B trunk+embedding) and graph copies (~10 GB) are co-resident, plus the retained zeros from (1). Measured: OOM at a 46 GB heap on the E2B export, in TraceToGraphBuilder.finalize.

3. Packed tensors cannot become graph constants at all — a keep-packed load traces them into opaque function arguments (we measured a func @gemma3n with 190+ weight args and zero dot ops), silently producing an unservable module. A loud error (or packed→dense dequant-on-extract) would be better than the silent signature leak; transformers now guards this with an arg-count check on its side.

Suggested directions (any one of the first two unblocks billion-scale export):

• Shape-only tensor data for VoidTensorOps (no backing array).
• Streaming/late constant extraction: dequant-or-copy per tensor at emission time, releasing each source afterwards, instead of building the full copied set up front.
• Reject packed frozen params in the tracer with a clear message.

Repro: SKaiNET-transformers :llm-inference:gemma3n:exportGemma3n on gemma-3n-E2B-it-Q4_K_M.gguf (48 GB host, -PexportMaxHeap=46g). The truncated GEMMA3N_LAYERS=4 export passes, isolating scale as the only variable.

[michalharakal]

Further characterization from the gemma3n export (layer-truncated run, GEMMA3N_LAYERS=4, so memory is NOT the limiter here — this is a distinct converter defect):

The emitted module carries only convert/transpose/reshape/slice ops and an empty return. The converter's own trail in the MLIR comments shows the chain breaking at the FIRST node and cascading:

// Conversion failed for node n0_gather: Gather operation requires 2 operands (weights, indices), got 1
// Conversion failed for node n1_reshape: Reshape operation requires exactly 1 operand, got 0
// Conversion failed for node n2_mulScalar: mulScalar requires exactly 1 tensor operand, got 0
…
// Conversion failed for node n1033_squeeze: Squeeze operation requires exactly 1 operand, got 0
// Warning: No output values found for return statement

plus one node that fails by name despite being in the converter's known-names list:

// Unsupported op 'weight' (type=constant) for node n1034_weight. Known names: [… 'weight' …]

So the embedding-gather's WEIGHT operand never links (the token embedding is dense FP32 and tied — the same tensor object serves token_embd.weight and output.weight), and every downstream node loses its operands transitively, leaving a compute-free module that still "succeeds". SmolLM2/FunctionGemma (also tied embeddings) emit fine at 135M/270M, so something about this trace shape — possibly the same frozen tensor object appearing as two differently-named params, or the weight-named constant node handling — breaks operand linkage.

Two asks beyond the memory items above:

1. Operand-linkage failures should be a hard error, not comments + empty return — transformers now arity-checks the emitted signature, but a compute-free module with exit 0 is a trap.
2. The n0_gather weight-linkage defect itself.

Repro: SKaiNET-transformers branch feat/gemma3n-iree-export, GEMMA3N_LAYERS=4 GEMMA3N_GGUF=…gemma-3n-E2B-it-Q4_K_M.gguf ./gradlew :llm-inference:gemma3n:exportGemma3n -PexportMaxHeap=40g.

[michalharakal]

Resolved on develop by #1248, #1249, #1250 (+#1251 stacked), #1253, and consumer-side SKaiNET-transformers#396.

Verification on the repro (gemma-3n-E2B-it-Q4_K_M.gguf, composite build of both repos' merged develop, -PexportMaxHeap=46g, full 30 layers):

	before	after
full E2B export at 46 GB	OOM in TraceToGraphBuilder.finalize	BUILD SUCCESSFUL in 42s
emitted module	convert/transpose/reshape only, empty return, exit 0	15,374 lines, 0 failure comments, 517 stablehlo.dot, 898 util.global weights, real return
function signature	190+ weight args	@gemma3n(%tokens: tensor<1x24xi32>, %per_layer_inputs: tensor<1x24x30x256xf32>)
weights	—	4629 MiB safetensors, 898 params, incl. the tied embedding (536,870,912 elems, 1 GiB bf16)

What each defect turned out to be:

1. Void zeros (perf(lang-core): shape-only void tracing — lazy placeholders instead of dense zeros (#1247) #1249): static-shape void ops now use lazy placeholders — a 4096×4096 traced stack tracks 0 bytes.
2. Double residency (perf(compile-dag)!: alias weight constants, fail loudly on packed frozen params (#1247) #1250): graph constants alias the live weight buffer instead of copyOf().
3. Packed → silent function args (perf(compile-dag)!: alias weight constants, fail loudly on packed frozen params (#1247) #1250): now PackedConstantException by default, PackedConstantHandling.DEQUANTIZE opt-in.
4. Soft failures / gather cascade (feat(compile-hlo)!: strict-by-default conversion errors, indexSelect routing, createBasic parity (#1247) #1248): ConversionErrorPolicy.STRICT default; resolveOperands throws MissingOperandException instead of mapNotNull-dropping (which also shifted operand positions — the "gather got 1 operand" cascade). indexSelect (the spelling the tracer emits) is now routable; createBasic registers NeuralNetOperationsConverter.
5. The "weight node fails despite being in known-names" mystery (feat(compile-hlo): array-free external constants via BufferHandle.Floats, Long-safe sizing (#1247) #1251): the tied embedding is 262144×2048 FP32 = Int.MAX_VALUE + 1 bytes — ByteArray(count * 4) overflowed, the NegativeArraySizeException landed in the generic catch, and the misleading known-names comment followed. External FP32 constants now ride an aliased FloatArray (BufferHandle.Floats) end-to-end with no byte serialization; sizes fold in Long; oversize serialization refuses with ConstantTooLargeException.
6. Strictness then surfaced one genuinely missing converter — clamp (feat(compile-hlo): lower clamp to stablehlo.clamp (#1247) #1253) — after which the export completes.