Home · ASPLOS 2024 · Machine Learning Systems

8-bit Transformer Inference and Fine-tuning for Edge Accelerators

Jeffrey Yu, Kartik Prabhu, Yonatan Urman, Robert M. Radway, Eric Han and Priyanka Raina

Conference
ASPLOS 2024
Published
2024
Pages
5-21
Citations
24

Running a Transformer's inference AND its fine-tuning on an edge chip means squeezing everything to 8 bits, but ordinary 8-bit float (FP8) spends its bits the same way at every magnitude while a Transformer's values cluster near 1.0 — so this Stanford work pits FP8 against 8-bit POSIT, a number format whose precision tapers to be sharpest exactly where the values live, quantizes not just the matrix multiplies but every operation (softmax, layernorm) by fusing them, adds 8-bit LoRA fine-tuning that cuts memory about 3x, and builds one accelerator that runs both formats — finding posit matches or beats FP8 on accuracy and, because its bit pattern makes softmax computable with a few bitwise tricks, yields a vector unit 33% smaller and 35% lower-power even though its matrix multiplier is slightly larger.

The argument in one glance

  1. Problem

    Transformers are too heavy to run — let alone fine-tune — on an edge chip

    Transformer models give state-of-the-art accuracy on language, speech, and vision, but they are enormous relative to an edge accelerator that has only kilobytes to a few megabytes of on-chip memory. Even MobileBERT, a deliberately small 25-million-parameter model, needs over 1.6 GB in half precision. And the edge does not just want inference: to keep data private and cut latency, devices increasingly need to FINE-TUNE a pre-trained model on their own downstream task locally. Inference and fine-tuning together, in a memory and power budget measured in milliwatts and megabytes, is the squeeze.

  2. What people did

    Drop to 8-bit — but only the matrix multiplies, and only in one float format

    The standard answer is 8-bit quantization. Prior work quantized the inputs to the big matrix multiplies (GEMMs) to 8 bits and left everything else — softmax, layernorm, the nonlinear activations, the accumulations — in high precision. And it largely assumed the 8-bit format would be a scaled-down float, FP8, with fixed field widths (e4m3 for the forward pass, e5m2 for the backward pass). Quantize the GEMM inputs, keep the format simple, call it done.

  3. Why it fell short

    Half-quantized, format-blind, and inference-only leaves accuracy and savings on the table

    Three gaps. FIRST, quantizing only GEMM inputs leaves the rest of the network in high precision, so much of the memory and energy is never actually saved — the softmax and layernorm traffic dominates on small models. SECOND, FP8 spends its precision uniformly across magnitudes, but a Transformer's weights and activations cluster near 1.0 — so a fixed-field float wastes resolution where the values are not and starves it where they are. THIRD, nobody had shown how to FINE-TUNE in 8 bits: activation gradients run one-to-two orders of magnitude past the 8-bit range and simply overflow, producing non-finite updates that destroy the model. An 8-bit story that touches only inference, only GEMMs, and only one format is missing most of the problem.

  4. This paper's idea

    Match the format to the data, quantize everything, fine-tune in 8 bits, and co-design the chip

    This Stanford work (Yu, Prabhu, Raina, and colleagues) does a full-stack study. It pits FP8 against 8-bit POSIT — a format whose variable-length fields make its precision TAPER, sharpest near 1.0 exactly where Transformer values live. It quantizes EVERY operation, not just GEMMs, by FUSING operations so intermediates stay high-precision inside a fused block and are quantized only at its boundary. It enables 8-bit FINE-TUNING via LoRA plus per-tensor scaling of the overflowing gradients, cutting fine-tuning memory about 3x. And it builds ONE accelerator that runs both formats, exploiting a bit-trick that computes softmax from a few posit bitwise operations. The payoff: posit matches or beats FP8 on accuracy at under 1% loss on large models, both formats cut accelerator area 30-34% and power 26-32% versus BFloat16, and posit's softmax trick makes its vector unit 33% smaller and 35% lower-power than FP8's.

Why this problem is worth solving

Putting a Transformer on an edge device is really a fight over bits. The models are huge, the chips are tiny, and the only way to close the gap is to represent each number in as few bits as accuracy allows. Eight bits is the aggressive target, and the question that decides everything is which 8-bit format to use. The default is

a scaled-down IEEE float packed into 8 bits, with FIXED field widths. The common variants are e4m3 (1 sign bit, 4 exponent bits, 3 mantissa bits), used for the forward pass, and e5m2 (5 exponent, 2 mantissa), used for the wider-range backward pass. Because the split between exponent and mantissa is fixed, an FP8 number has the SAME relative precision at every magnitude — the same number of mantissa bits whether the value is near 0.001 or near 1000. That uniformity is simple to build and simple to reason about, but it means the format cannot concentrate its scarce precision anywhere in particular
. The alternative this paper takes seriously is a different way of spending those same eight bits.

The reason format choice matters so much is that a Transformer does not produce numbers uniformly across all magnitudes — its weights, activations, and attention scores cluster near 1.0. A format that puts its sharpest resolution there, and lets precision fall off for the rare very-large or very-small values, fits the data better than one that spreads precision evenly. That is exactly what a

an 8-bit format with VARIABLE-length fields: sign, then a run-length “regime” field (2 to 7 bits), then maybe an exponent bit, then whatever fraction bits are left. The regime encodes a coarse power of two by its run length, and the number of fraction bits that survive depends on how big the regime is — so values near 1.0 get many fraction bits (high precision) while values far from 1.0 get few (low precision). This TAPERED precision is sharpest right where a neural network’s values live, and near 1.0 a Posit8 carries roughly a full extra decimal digit of accuracy over the comparable FP8. The cost is that decoding the variable fields needs real hardware a fixed-field float does not
does. So the study is not academic: which 8-bit format wins decides accuracy, and — because the two formats need different decoders and enable different shortcuts — it also decides the silicon. And underneath the format question sit two more the prior work skipped: quantizing the whole network rather than only its matrix multiplies, and fine-tuning (not just running) the model in 8 bits without the gradients overflowing. Getting all three right is what turns 8-bit from a partial trick into a way to actually fit a trainable Transformer on an edge chip.

What people did before, and where it broke

  • Quantize only the GEMM inputs (most 8-bit inference work). Cast the inputs to the big matrix multiplies to 8 bits, leave softmax, layernorm, the nonlinearities, and the accumulations in high precision. Real speedup on the matmuls, but the un-quantized operations still dominate memory and energy on small models — the savings are left half-collected.
  • Assume FP8 is the format (float-centric quantization). Treat 8-bit as a shrunken float with fixed exponent/mantissa fields. Simple hardware, but the fixed split spreads precision uniformly across magnitudes while the model’s values pile up near 1.0, so FP8 wastes resolution where there is no data and lacks it where there is — especially on models with wide activation spreads.
  • Inference-only 8-bit (the bulk of the literature). Optimize the forward pass and stop. But edge devices need to fine-tune locally for privacy and latency, and nobody had shown stable 8-bit fine-tuning: activation gradients exceed the 8-bit range by one-to-two orders of magnitude and overflow to non-finite values, breaking training.
  • Full-precision softmax in hardware (accelerator norm). Compute attention’s softmax (exponential, reciprocal) in high precision because it is numerically delicate. Correct, but the vector unit that does softmax then dominates area and power in an attention-heavy model, and none of the 8-bit format’s structure is used to make it cheaper.
  • The gap. No prior work jointly chose the RIGHT 8-bit format for Transformer data, quantized EVERY operation, enabled 8-bit FINE-TUNING, AND co-designed one accelerator whose hardware exploits the format to make softmax nearly free. This paper does all four, and measures the trade between the two formats across the whole chip.

How this work makes an 8-bit edge Transformer

Pick the format, quantize everything, fine-tune in 8 bits, and build the chip around it

  1. Match the format to the data distribution new

    Compare FP8 (fixed e4m3/e5m2 fields) against Posit8 (tapered precision from variable regime/exponent/fraction fields). Because a Transformer's values cluster near 1.0, posit's sharpest-near-1.0 precision fits them better — carrying about one extra decimal digit of accuracy there than the comparable FP8. Neither wins everywhere: FP8 does better on small models whose activations spread wide, so the format is selected per model.

  2. Quantize EVERY operation via fusion new

    Prior work quantized only GEMM inputs; this quantizes softmax, layernorm, and the nonlinearities too. The mechanism is OPERATION FUSION: a chain of operations is executed as one fused block that keeps its intermediates in high precision internally and quantizes to 8 bits only at the block boundary — so the whole network runs 8-bit without accumulating rounding error at every step. Small models need more aggressive fusion to stay under 1% loss.

  3. 8-bit fine-tuning with LoRA and gradient scaling new

    Fine-tune by freezing the large pre-trained weights and training a small low-rank LoRA adapter in 8-bit arithmetic — so only the adapter's activations and gradients need storing. Activation gradients overflow the 8-bit range by one-to-two orders of magnitude, so per-tensor scaling rescales them into range before quantizing. Combined, this cuts fine-tuning memory about 3x (roughly 500 MB down to 150 MB on MobileBERT).

  4. One accelerator, both formats new

    A single design: an NxN systolic array (the matrix unit) for GEMMs plus an N-lane vector unit for element-wise operations and softmax. The 8-bit MACs are far smaller than a BFloat16 MAC either way; posit needs extra encode/decode logic for its variable fields, making its MAC slightly larger than FP8's — a cost that has to be paid back elsewhere.

  5. Approximate softmax from posit bit-tricks new

    The payback: posit's bit layout makes the transcendental softmax functions almost free. A sigmoid is approximated by flipping the most-significant bit and shifting right; the exponential is built as a reciprocal of that sigmoid; the reciprocal is a cheap bitwise operation on the posit encoding. The one hazard is that the posit exp approximation does not go to 0 for large-negative inputs, which would break attention masking — fixed by thresholding at about -3, cutting a MobileBERT accuracy loss from 9.8% down to 0.3%.

  6. Account for the cost across the WHOLE chip new

    Judge the format on the entire accelerator, not one unit. Both 8-bit formats cut area 30-34% and power 26-32% versus BFloat16. Posit's matrix MAC is slightly larger than FP8's, but its softmax trick makes its VECTOR unit 33% smaller and 35% lower-power than FP8's — so for an attention-heavy Transformer, the format that costs a little more on the matmul pays for itself on the softmax.

Evaluated at 40nm, 200 MHz, with a 32x32 systolic array, across BERT (base/large), MobileBERT, RoBERTa, Whisper, GPT-2, and LLaMA 2 (7B/13B), on SQuAD, GLUE, LibriSpeech, and WikiText-103, against BFloat16 and FP8 baselines. Fully open-sourced; earned all three ASPLOS artifact badges.

The idea

Spend the same eight bits where the data actually is. FP8 gives every number the same relative precision at every magnitude, which sounds fair but is wasteful when the data is not spread evenly. A Transformer’s weights, activations, and attention scores pile up near 1.0, and posit’s tapered precision is sharpest exactly there — many fraction bits near 1.0, few out in the rare tails. So posit fits the distribution the model produces, buying about an extra decimal digit of accuracy near 1.0 for free, and the measurements bear it out: Posit8 matches or beats FP8 across most models at under 1% loss on the large ones. This is the match-the-representation-to-the- data discipline the site keeps meeting — the win comes not from more bits but from putting the bits you have where the values live.

Do not judge a number format by the matrix unit alone. The tempting comparison is FP8 versus posit on the GEMM MAC, and there FP8 wins narrowly — posit’s variable fields need decode logic that makes its multiplier slightly larger. But that is the wrong accounting boundary. An attention-heavy Transformer spends a large fraction of its vector-unit time in softmax, and posit’s bit structure lets softmax’s exponential and reciprocal be approximated with a handful of bitwise operations, so posit’s vector unit comes out 33% smaller and 35% lower-power than FP8’s. Charge the format for what it costs across the WHOLE chip and the verdict flips: posit pays a little on the matmul and wins more on the softmax. It is the same true-cost-over-a-convenient-proxy lesson the site finds again and again — the proxy (MAC area) said FP8, the true cost (whole-accelerator area and power on real attention) said posit.

Quantize the parts everyone else left alone, and fine-tune too. The prior recipe quantized GEMM inputs and stopped, so the un-quantized softmax and layernorm quietly kept dominating small-model cost — an Amdahl trap where the un-optimized fraction caps the win. This work quantizes EVERY operation by fusing chains so intermediates stay precise internally and only the boundaries round to 8 bits. And it crosses the line prior work would not: fine-tuning in 8 bits. The obstacle is that activation gradients overflow the 8-bit range by one-to-two orders of magnitude, so it rescales them per-tensor before quantizing and freezes the big weights while training only a low-rank LoRA adapter — cutting fine-tuning memory about 3x and keeping accuracy within half a percent on the models tested. Attack the operations others ignored, and the trainable edge Transformer becomes real.

Choosing bits for one Transformer, end to end

From the data distribution to the silicon

Walking the design decisions on an attention layer, with numbers from the paper's evaluation at 40nm.

Step 1 of 5

Look at where the numbers are

The weights and activations flowing through an attention layer are not spread evenly across magnitudes — they cluster near 1.0, roughly Gaussian. An 8-bit float (FP8) gives the same precision at 0.001 and at 1000, so most of its resolution sits where there is no data. An 8-bit posit tapers: near 1.0 it keeps many fraction bits, far away it keeps few. Matching the format to the distribution is the first and cheapest win.

near 1.0 a Posit8 carries about one extra decimal digit of accuracy over the comparable FP8
posit range with es=1 spans about 2^-12 to 2^12 — enough for weights, activations, and gradients

Why it works

  • A Transformer’s values cluster near 1.0 while FP8 spreads precision uniformly → posit’s tapered precision puts its sharpest resolution near 1.0 → about an extra decimal digit of accuracy where the data lives, at the same 8 bits.
  • Quantizing only GEMM inputs leaves softmax and layernorm in high precision (an Amdahl cap) → operation fusion quantizes every operation while keeping intermediates precise inside each fused block → the whole network’s memory and energy are actually saved, not half of them.
  • Posit’s matrix MAC is slightly larger than FP8’s, so the matmul favors FP8 → but posit’s bit layout makes softmax’s exp and reciprocal a few bitwise operations → its vector unit is 33% smaller and 35% lower-power, and the whole-chip verdict flips to posit on attention-heavy models.
  • Activation gradients overflow the 8-bit range by one-to-two orders of magnitude → per-tensor scaling rescales each tensor into range before quantizing, and LoRA trains only a small adapter → 8-bit fine-tuning becomes stable and cuts memory about 3x.
  • The approximate posit exp fails to decay to 0 for large-negative scores, leaking masked attention positions → thresholding the exponent at about -3 forces those to zero → a would-be 9.8% MobileBERT accuracy loss shrinks to 0.3%.

What it costs

  • No single format wins. Posit is better on large, wide-dynamic-range models; FP8 is better on small models whose activations spread wide (MobileBERT with stacked feed-forward layers). The system must SELECT the format per model, which means supporting both in hardware.
  • Small models are fragile. Staying under 1% loss on small models (MobileBERT tiny, GPT-2 small) needs the most aggressive operation fusion and per-tensor scaling; without them the loss exceeds 1%. Larger models (1.5B and up) are far more robust to the fusion level.
  • The softmax approximation is delicate. The posit exp does not converge to 0 as inputs go to negative infinity, so attention masking needs an empirically-tuned threshold (about -3) and a custom gradient derivation for the backward pass — a hand-tuned patch, not a clean function.
  • Low-precision fine-tuning is unstable. Even with scaling, 8-bit training can produce non-finite gradients; the paper relies on lower learning rates, checkpoint resumption, and fusing the final classification layer to keep runs alive.
  • Posit costs extra silicon per MAC. Decoding posit’s variable-length fields needs encoder/decoder logic FP8 does not, making the posit MAC slightly larger; only the vector-unit savings pay it back, so the win is conditional on the workload being softmax-heavy.

The evaluation

Two 8-bit formats (FP8 in e4m3/e5m2, and Posit8 with es of 1 or 2) across BERT base/large, MobileBERT, RoBERTa, Whisper, GPT-2, and LLaMA 2 (7B and 13B), on SQuAD, GLUE, LibriSpeech, and WikiText-103, with a 32x32 systolic-array accelerator synthesized at 40nm, 200 MHz, compared against BFloat16 and across the two 8-bit formats.

Three ways to spend the bits on an edge Transformer

PropertyBFloat16 baseline8-bit float (FP8)8-bit posit (this work)
Precision vs. data distribution
What gets quantizedNothing (full precision reference)
Softmax in hardware
Accuracy (large models, under 1% loss)Reference (e.g. BERT-large 93.2 F1)
Chip area / power vs BF16
Fine-tuning (not shown above): LoRA plus 8-bit cuts fine-tuning memory about 3x (roughly 500 MB to 150 MB on MobileBERT) while staying within half a percent of BFloat16 on GLUE and SQuAD. Both 8-bit MACs are far smaller than BF16's (e.g. Posit8 and E4M3 MACs are roughly half BF16's area and well under half its power).Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. FP8 and 8-bit posit both use exactly 8 bits, yet posit is more accurate on Transformer data. Predict why the SAME number of bits buys posit more accuracy where it counts.

  2. 2. On the raw matrix-multiply MAC, FP8 is slightly smaller and lower-power than posit — yet the paper argues posit is the better accelerator format for attention-heavy Transformers. Predict how both can be true.

  3. 3. Prior 8-bit work quantized the inputs to the matrix multiplies and left softmax and layernorm in high precision. Predict why this leaves much of the potential savings uncollected, especially on small models.

  4. 4. To fine-tune in 8 bits, the paper rescales activation gradients per-tensor and trains only a small LoRA adapter, rather than widening the number format. Predict why this is the right response to gradients that overflow the 8-bit range by one-to-two orders of magnitude.

Sources

Related papers on Machine Learning Systems

See all 212 →