Skip to content

granite

granite

Granite (GraniteForCausalLM, IBM Granite 3.x / 4.x dense).

Llama's tree and names, with four scalars from the config applied around it:

  • embedding_multiplier: the model multiplies the embedding module's output by it before the first block, so token_embeddings (the module's output) times the multiplier is layers[0].input.
  • attention_multiplier: the attention's softmax scale in place of 1/sqrt(head_dim), passed to the shared interface; nothing to override.
  • residual_multiplier: the block adds h * residual_multiplier for each sublayer, not h. What reaches the residual stream is the scaled term, so attention_output and mlp_output are the module's output times the multiplier. That product is a computed copy, so each value divides on assignment, and a transform divides an edited copy back into the module's output so in-place edits reach the model.
  • logits_scaling: the model divides the head's output by it, so the family defines project_on_vocab with that step and the logit lens on the last block equals logits.

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm'}

ENVOYS module-attribute

ENVOYS = {GraniteDecoderLayer: Layer, GraniteAttention: Attention, GraniteMLP: Mlp}

residual_multiplier

residual_multiplier(envoy: Envoy) -> float

What the block multiplies each sublayer's output by before adding it.

scaled_back

scaled_back(envoy: Envoy, edited: Tensor, raw) -> object

The module output an edited scaled copy stands for: edited / residual_multiplier, rewrapped.

An unedited copy hands the module's own output back untouched, so a read leaves the forward bit-identical even where dividing by the multiplier would round.

project_on_vocab

project_on_vocab(model: StandardizedTransformer, hidden: Tensor) -> Tensor

The logit lens as the model makes its logits: the final norm, lm_head, then divided by logits_scaling.