distilgpt2-attn-onnx

distilgpt2-attn-onnx — distilgpt2 re-exported so you can actually read its attention

An ONNX export of distilgpt2 whose graph emits real attention probabilities as named outputs (attn_0 … attn_5).

The stock ONNX exports on the Hub emit logits + present.* (KV cache) and no attention tensors — they simply are not in the graph. If you want to visualise what a real transformer's heads are doing in a browser, there is nothing to read. This export fixes exactly that, and nothing else.

Built for the attention lab in ai-cowboys-learn, which previously had to fall back to hand-built reference patterns and say so on screen. It doesn't anymore.

⚠️ Do not quantize this to int8. It fails silently.

This is the most important line in this card.

Int8 dynamic quantization corrupts the attention weights — up to 0.34–0.60 absolute error on a probability. Measured, not estimated:

And here is the trap: the int8 output still passes every self-consistency check you would naturally reach for. Every row still sums to 1.0. The causal mask still holds exactly. The heatmap still looks like attention. It is a softmax of wrong logits — internally consistent and completely false. Sanity checks that only interrogate the output cannot catch this. The only thing that catches it is comparing against a reference implementation.

Errors compound through the layers and get amplified by the softmax exponential, so it fails with only the Gemms quantized too — it isn't the embedding's fault.

Use fp16 or fp32. Never q8/int8.

Note for transformers.js users: DEFAULT_DEVICE_DTYPE_MAPPING[wasm] = 'q8', so omitting dtype silently gives you the corrupt variant. Always pass it explicitly.

How it was exported

Two things matter, and both are easy to get wrong:

1. attn_implementation="eager" is mandatory. The default sdpa (and flash) return attentions=None silently — no error, no warning. Your export "succeeds" and emits nothing. This will cost you a day if you don't know it.
2. The HF ONNX exporter does not expose attentions. You have to wrap the model in an nn.Module whose forward returns the attentions as real outputs, then export that.

model = AutoModelForCausalLM.from_pretrained(
    "distilgpt2", attn_implementation="eager", dtype=torch.float32)   # eager!

class AttnOnly(nn.Module):
    def __init__(self, m):
        super().__init__()
        self.trunk = m.transformer          # no lm_head
    def forward(self, input_ids, attention_mask):
        o = self.trunk(input_ids=input_ids, attention_mask=attention_mask,
                       output_attentions=True, use_cache=False, return_dict=True)
        return tuple(o.attentions)          # -> attn_0 .. attn_5

torch.onnx.export(
    AttnOnly(model).eval(), (ids, mask), "model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=[f"attn_{i}" for i in range(6)],
    dynamic_axes={"input_ids": {0: "batch", 1: "sequence"},
                  "attention_mask": {0: "batch", 1: "sequence"},
                  **{f"attn_{i}": {0: "batch", 2: "sequence", 3: "sequence"} for i in range(6)}},
    opset_version=14, do_constant_folding=True)

No logits — this reads attention, it does not generate text

The lm_head is dropped. There is no logits output. This model cannot generate text, and that is deliberate:

• distilgpt2 ties lm_head to wte, and the export materialised the embedding twice — wte.weight (154 MB) plus a transposed copy onnx::MatMul_1629 (154 MB). 308 MB of a 482 MB export was the same matrix stored twice. Dropping the head lets DCE remove it: 482 MB → 306 MB.
• AutoModelForCausalLM in transformers.js loads it fine anyway (verified in-browser).

There is also no KV cache (use_cache=False), so no present.* outputs — one clean forward pass over the whole sequence, which is what you want for visualisation.

Verification

Every number here was measured, not assumed. Checked at four sequence lengths (2, 10, 13, 17) — deliberately including lengths ≠ the 10-token trace length, because the tracer warns it may bake the mask in as a constant. It does not; dynamic axes hold.

For fp32, on every sentence:

• max abs diff vs PyTorch output_attentions=True: 1.669e-06 (mean 4.6e-08)
• shapes: [batch, 12, seq, seq] × 6 layers (distilgpt2 = 6 layers × 12 heads)
• softmax rows sum to 1.0: max |rowsum − 1| = 2.4e-07
• causal mask exact: max upper-triangle = 0.000e+00

Verified end-to-end in headless Chromium via transformers.js, comparing tensors pulled out of the browser against PyTorch:

fp16 on WASM is fine. It is WebGPU that is picky about fp16 — the opposite of what you might expect.

As a real check that the weights mean something: layer 2, head 11 is a previous-token head (mean attention to position i−1 = 0.951; rows peak at i−1 at ≈1.000) on "The cat sat on the mat because it was warm". That was measured from this export in a browser, not asserted from the literature.

Usage (transformers.js)

import { AutoTokenizer, AutoModelForCausalLM } from '@huggingface/transformers';

const id = 'The-AI-Cowboys-Projects/distilgpt2-attn-onnx';
const tokenizer = await AutoTokenizer.from_pretrained(id);
const model = await AutoModelForCausalLM.from_pretrained(id, {
  dtype: 'fp16',    // or 'fp32'. NEVER 'q8' — see the warning above.
  device: 'wasm',   // fp16 throws on webgpu adapters without shader-f16
});

const enc = await tokenizer('The cat sat on the mat because it was warm');
const out = await model({ ...enc, output_attentions: true });

// out.attn_0 … out.attn_5 — each [1, 12, seq, seq], row-stochastic, causally masked
console.log(out.attn_0.dims);   // [1, 12, 10, 10]

Files

transformers.js selects the file from dtype: fp32 → model.onnx, fp16 → model_fp16.onnx.

License & attribution

Apache-2.0, inherited from distilgpt2 (Hugging Face), itself a distilled version of OpenAI's GPT-2. This is a derivative work — a re-export, with no retraining and no weight modification beyond the fp16 cast — and is released under Apache-2.0 as well. All model behaviour, biases, and limitations of distilgpt2 carry over unchanged; see the distilgpt2 card for those.

Exported for ai-cowboys-learn.