Problem
Gemma 3n models (e.g. yujiepan/gemma-3n-tiny-random) fail to load with StandardizedVLM / load_model():
ValueError: layers_input[0] has shape torch.Size([4, 1, 3, 32]) != (1, 3, 32) in yujiepan/gemma-3n-tiny-random architecture.
Root cause
Gemma 3n uses AltUp (Alternating Updates), a technique from this NeurIPS 2023 paper that fundamentally changes the hidden state shape flowing through decoder layers.
Before any decoder layer, the token embedding is projected into altup_num_inputs (=4) copies and stacked into a 4D tensor (altup_num_inputs, batch, seq, hidden). All decoder layers receive and return this 4D tensor. Only the copy at index 0 (the "active" one) goes through attention and MLP; the others are updated via learned predict/correct steps. After all layers, the 4D tensor is collapsed back to 3D.
Why this is hard to support
nnterp assumes hidden states are 3D (batch, seq, hidden) throughout. Supporting AltUp would require changes across most of the codebase:
layers_input[i] / layers_output[i]: would need to either return the full 4D tensor (breaking the standard interface) or only the active copy at index 0 (lossy)logit_lens: the "hidden state" at each layer is really 4 hidden states. The active one (index 0) is most meaningful, but the others carry information too- Steering: modifying
layers_output[i] would need to account for all 4 copies, not just one
- Shape assertions:
check_io in rename_utils.py rejects the 4D shape
This is a genuinely new architecture pattern, not a missing rename config or nnsight issue.
Status
Marking as won't fix for now unless there's demand. If you're trying to use nnterp with Gemma 3n and need this, please comment here.
🤖 Generated with Claude Code
Problem
Gemma 3n models (e.g.
yujiepan/gemma-3n-tiny-random) fail to load withStandardizedVLM/load_model():Root cause
Gemma 3n uses AltUp (Alternating Updates), a technique from this NeurIPS 2023 paper that fundamentally changes the hidden state shape flowing through decoder layers.
Before any decoder layer, the token embedding is projected into
altup_num_inputs(=4) copies and stacked into a 4D tensor(altup_num_inputs, batch, seq, hidden). All decoder layers receive and return this 4D tensor. Only the copy at index 0 (the "active" one) goes through attention and MLP; the others are updated via learned predict/correct steps. After all layers, the 4D tensor is collapsed back to 3D.Why this is hard to support
nnterp assumes hidden states are 3D
(batch, seq, hidden)throughout. Supporting AltUp would require changes across most of the codebase:layers_input[i]/layers_output[i]: would need to either return the full 4D tensor (breaking the standard interface) or only the active copy at index 0 (lossy)logit_lens: the "hidden state" at each layer is really 4 hidden states. The active one (index 0) is most meaningful, but the others carry information toolayers_output[i]would need to account for all 4 copies, not just onecheck_ioinrename_utils.pyrejects the 4D shapeThis is a genuinely new architecture pattern, not a missing rename config or nnsight issue.
Status
Marking as won't fix for now unless there's demand. If you're trying to use nnterp with Gemma 3n and need this, please comment here.
🤖 Generated with Claude Code