OLMo-2: reordered (post) normalization and a loader-clobber root-cause

1. Contribution

Claimed. An independent dense-inference-engine port of OLMo-2-1B (AllenAI), implemented without a vendor BLAS dependency, reproduces a recent Hugging Face transformers f16 prefill-logit reference at min cosine 0.999998, argmax 12/12 on an RDNA2-class AMD GPU, given identical token ids. The port correctly implements OLMo-2's three architectural departures from the standard pre-norm Llama-style block: (a) reordered / post-normalization — there is no input normalization; the raw residual feeds the attention and MLP projections, and the post-attention and post-feedforward normalizations are applied to the sublayer OUTPUTS before they are added back to the residual; (b) full-dimension QK-normalization (RMSNorm over the whole Q/K projection, NOT per-head). A second, engineering-level contribution is the identification and fix of a loader-clobber root cause: the engine's generic weight loader overwrote the MLP normalization with the post-attention normalization, corrupting the forward pass until a model-specific guard was introduced.

NOT claimed. This is a correctness-validation result, not a novel algorithm and not a benchmark-scale evaluation. It does NOT claim generation quality, decode/throughput performance, or any wall-clock figure. It does NOT claim parity for larger OLMo-2 variants, nor that the int8 path (measured cosine 0.998700) matches f16. The reordered-norm and QK-norm mechanisms are OLMo-2's own design, not contributions of this work; the contribution is the faithful independent reproduction plus the root-cause fix.

2. Methods

3. Results

4. Ablations / Negative controls

The source records the loader-clobber bug as a before/after observation rather than a controlled ablation: with the generic loader overwriting the MLP normalization, the forward pass is wrong; after introducing the model-specific guard, parity reaches min cosine 0.999998 / argmax 12/12. This is the closest thing to a negative control available, but it is a bug-fix delta, not a systematically swept ablation of each architectural feature (reordered-norm, full-dim QK-norm) in isolation. No per-feature force-off ablation is recorded.

5. Threats to validity / Limitations

• Sample size. Parity covers 12 token positions (a single small set of prompt positions). This is a correctness smoke test, not a benchmark-scale evaluation. A high cosine across 12 positions does not establish behavior across long contexts, diverse prompts, or rare token paths.
• Parity is not quality. Logit cosine and argmax match against an oracle measure the numerical faithfulness of one forward pass. They say nothing about end-to-end generation quality, coherence, or task accuracy.
• No throughput / single forward. No decode rate, latency, or variance is measured. There is no multi-run statistic; the reported numbers are point measurements.
• Precision ceiling. The engine's highest precision here is f16; there is no f32 reference forward inside the engine, so the irreducible f16 ULP noise floor bounds achievable cosine. The int8 path (0.998700) is reported but explicitly NOT claimed to match f16.
• Scope. Claims are bounded to OLMo-2-1B on the tested RDNA2-class GPU versus a recent transformers release. Larger OLMo-2 variants, other hardware, and other reference versions are untested.
• Ablation rigor. The root-cause fix is validated as a before/after bug delta, not as an isolated per-feature ablation; the load-bearing role of reordered-norm vs full-dim QK-norm individually is not separately demonstrated.
• Oracle trust. Results assume the reference release loads OLMo-2 weights correctly; the oracle was not independently re-validated beyond version pinning.
• evidenceLevel. Retained as "experimental"; given the 12-position sample this is a correctness-validation experiment, not a broad empirical study — readers should weight it accordingly.

6. Reproducibility / Provenance

• Oracle pin: a recent Hugging Face transformers release (specific version pinned internally).
• Hardware: RDNA2-class AMD GPU (16 GB).
• Fix locus (high level): the generic weight loader is guarded by a model-specific check so OLMo-2's reordered normalizations load into the correct slots and the MLP normalization is no longer overwritten.
• Independence check: the runtime links only the AMD HIP runtime plus system libraries — no vendor BLAS / DNN math library.
• Validate the oracle first: confirm the pinned reference actually loads the OLMo-2 weights (not random-init) before trusting parity numbers.
• Build / run: compiled with a standard HIP compiler toolchain targeting the RDNA2 architecture; exact harness invocation not recorded in source notes.