zaya¶
ZAYA (ZayaForCausalLM, Zyphra's ZAYA1).
Llama's containers over a block with residual scaling modules: after each
sublayer a ZayaResidualScaling merges the sublayer's output o into the
stream r with learned per-channel scales and biases::
merge(o, r) = (o + hidden_states_bias) * hidden_states_scale + (r + residual_bias) * residual_scale
h = post_attention_residual_scale(self_attn(input_layernorm(x)), x)
out = post_mlp_residual_scale(mlp(post_attention_layernorm(h)), h)
What each sublayer adds is its scaled term, (o + hidden_states_bias) *
hidden_states_scale, so attention_output and mlp_output are that
product: a computed copy of the module's output, inverted on assignment, with a
transform carrying an in-place edit back, as on Granite. The stream itself is
also rescaled, so the contribution identity is
layer_output == (h + residual_bias) * residual_scale + mlp_output with
h == (input + residual_bias) * residual_scale + attention_output (each merge's
own parameters, layers[i].post_*_residual_scale._module, which each sublayer
reads through its parent as merge), the plain one only
at their initial values (scales one, biases zero). The model also scales and
shifts the embedding module's output before the first block.
The attention runs the shared eager forward; its queries and keys come from a
compressed convolutional projection (qkv_proj, a causal convolution over q/k
and a delayed half of the values) and qk_norm, so there are no q_proj /
k_proj on the module. The block returns (hidden_states,
prev_router_hidden_states): the mixture's router carries a state from one block
to the next. The mixture (ZayaSparseMoeBlock) is on every block and returns
(hidden_states, router_state); its experts are moe_intermediate_size wide,
which the root's intermediate_size reports (the config has no dense width).
qk_norm L2-normalizes queries and keys and multiplies the keys by a learned
per-head temp; at its initial value of zero every key is zero and the pattern
uniform.
The mixture is a Moe; its router, gate (aliased router), mixes the
previous block's router state into its input and scores num_experts + 1
classes with a small MLP (router_mlp), whose output is router_logits: the
last column is skip. A slot that picks it runs no expert: its weight is 0 and
its index 0, an alias of expert 0, so usage counts mask expert_weights == 0.
The router picks on its softmax plus balancing_biases; on ZAYA1-8B and
ZAYA1-74B-preview (num_experts_per_tok 1) the skip class's bias is -1 and some
expert's is positive on every block, so skip is never chosen there; the tiny picks it.
module-attribute
¶
RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'gate': 'router'}
module-attribute
¶
ENVOYS = {ZayaDecoderLayer: Layer, ZayaAttention: Attention, ZayaSparseMoeBlock: Mlp}
Bases: Moe
ZAYA's mixture of experts ((hidden_states, router_state)); the block adds its output shifted and scaled by post_mlp_residual_scale.
router_logits are router.router_mlp's output, num_experts + 1
columns (the last is skip), already [batch, seq, ...].
property
¶
The envoy of the merge that follows this sublayer, its sibling in the block.
intermediate_size(model: StandardizedTransformer) -> int
The MLP width is one expert's, moe_intermediate_size: every block is a mixture and the config has no dense width.