MLPSpeculatorPreTrainedModel

MLPSpeculatorPreTrainedModel at 4 prediction stages and stage width 4096. 2 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

No forward was recorded for this shape; the block is read from the published configuration and from MLPSpeculator in fms-extras at commit 16339f7.

From granite-3.0-8b-instruct-accelerator

# A draft head, not a model: it proposes the next 4 tokens for the base
# model to verify. No attention, no positions, no cache: each stage sees
# one state vector and one token id.

# state = the base model's last hidden row      [4096]
# token = the id the base model just emitted

state = rmsnorm(state) / sqrt(2)          # `scale_input: true`: ln0 has no
                                          #   weight and no bias, and the
                                          #   / sqrt(2) sits outside it
state_weight = 0.5 ** (0.5 / 4)   # = 0.917004
emb_weight   = sqrt((1 - state_weight**2) * 4096 / 2)   # = 18.0512
alpha        = emb_weight / state_weight                # = 19.6849
# Derived from n_predict and inner_dim; neither is a tensor. Only the
# ratio is applied: the norm below absorbs the missing state_weight,
# up to the norm's eps (scaling x by c turns eps into eps / c**2).

for i in 0 .. 3:
    state = proj[i] @ state               # [4096]; proj[0] reads 4096, every
                                          #   stage after it is square
    state = state + alpha * emb[i][token]
    state = rmsnorm(state, ln[i].weight, ln[i].bias, eps=1e-06)
    #  ^ `LayerNormParameterized` with its mean subtraction off:
    #    x * rsqrt(mean(x^2) + eps), then weight * x + bias.
    state = gelu(state)                   # exact (erf) GELU, not tanh
    logits = head[i] @ state              # [49155]
    token  = argmax(logits)               # greedy; see Notes for the
                                          #   candidate tree
    draft.append(token)

Differs from the block above

The same arithmetic with these values changed. Everything not listed is as above.

emb_dim and inner_dim can differ. On granite-20b-code-instruct-accelerator, emb_dim (6,144) is the base model's hidden width, which proj[0] reads, and inner_dim (4,096) is the head's own width, which every later stage runs at (speculator.py L62-L67).

Notes

This head drafts 4 tokens ahead, one per stage, each stage feeding the next. It verifies nothing: in fms-extras' greedy speculative_generate, the base model keeps the longest prefix of a candidate that matches its own next tokens (generation.py L242-L244). A wrong draft therefore costs accepted tokens per step, not correctness.

The norm is an RMS norm with a bias, not a LayerNorm. MLPSpeculator builds each ln[i] as LayerNormParameterized(inner_dim, elementwise_shift=True, elementwise_scale=True) (speculator.py L71-L78), whose defaults are use_mean=False, eps=1e-06 (fms layernorm.py L27-L34, L59-L72 at v0.0.6). Subtracting the mean would load and run, and draft different tokens. The activation is nn.GELU(), PyTorch's exact form (L86).

state_weight and emb_weight are computed, not stored (speculator.py L83-L85), and only their ratio is applied: upstream writes torch.add(state, z, alpha=self.emb_weight / self.state_weight) and relies on the following norm (L162-L165). A port that scales both terms separately differs only through that norm's epsilon, but it is not the published arithmetic.

scale_input: true (granite-3.0-8b-instruct-accelerator) adds a parameter-free norm on the incoming state. ln0 is built with elementwise_shift=False, elementwise_scale=False (speculator.py L79-L82), so the checkpoint has no tensor for it, and the inference path divides its output by sqrt(2) (L156-L157).

tie_weights: true (granite-3.0-8b-instruct-accelerator) shares tensors across stages. emb[i], head[i] and ln[i] all point at stage 0's tensors and proj[i] for i ≥ 2 at proj[1] (speculator.py L88-L106). A loader has to resolve those aliases rather than expect one tensor per stage.

top_k_tokens_per_head: [4, 3, 2, 2] with n_candidates: 4 configures generate_suffixes, which keeps the top k tokens at each stage and returns the 4 most probable candidate sequences (speculator.py L116-L187). Greedy argmax at every stage is one path through that tree.

Geometry

Weight structure

Both checkpoints of this shape publish only pytorch_model.bin, whose tensor list cannot be read without downloading the file, so the tensor names and counts this sheet implies cannot be checked against the published headers.

Other shapes of MLPSpeculatorPreTrainedModel

Only the columns that differ are shown; a cell with several values means the checkpoints of that shape disagree.

How this was checked

Status and evidence grade in the assessment (see what the grades mean):

• each of the 2 checkpoints of this shape: runs, own-pack — a stored copy of this checkpoint exists; nothing recorded from it

Parts of the author's unpublished implementation that exist for this shape:

• a draft head implemented from the published configuration
• a checkpoint loader, which is where published tensor names are read
• an independent reference implementation of this architecture in Python, driving the published modelling code — an executable statement of what the model should compute, written against the publication rather than against any one implementation of it

Checkpoints with this architecture