granite_swa¶
Granite SWA (GraniteSWAForCausalLM).
Granite's tree, block and multipliers (granite.py: scaled residual adds,
embedding_multiplier, logits_scaling), with layer_types mixing
sliding-window and full-attention blocks, a rotary embedding per block (or none,
NoPE, where the model passes no position embeddings), and a learned attention
sink per head. The sink does not join the softmax as a column (as on GPT-OSS):
the eager forward takes the softmax over the real keys, mixes the values with it,
and then scales each query's head output by sigmoid(logsumexp(scores) - sink),
the share of the mass a softmax with the sink as one more logit would leave on the
real keys. So attention_probabilities is that softmax (F_dropout_0), whose
rows sum to one, attention_scores its input, and attention_head_outputs
the interface's output, after the sink scale. The eager forward spells its
softmax and dropout F.softmax / F.dropout, so both are redefined on those
operations.
module-attribute
¶
module-attribute
¶
Bases: Attention
Granite SWA's attention: Granite's scaled contribution, and a sink applied to the head outputs after the softmax.
The pattern is the softmax over the real keys, rows summing to one; the sink's share scales the head outputs, not the pattern.
Bases: Mlp
Granite SWA's MLP: Granite's, the block adds its output times residual_multiplier.