Research map / Architectures and efficiency / Quantization
Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Jetfire, an INT8 pretraining method for transformers that keeps INT8 data flow across linear and non-linear operators with per-block (32x32) quantization, matching FP16 accuracy and giving 1.42x end-to-end block speedup on RTX 4090.
Problem
Fully quantized training of transformers is limited by quantize-compute-dequantize overheads, memory-bound non-linear operators, and accuracy loss from channel-wise outliers under per-token or per-tensor INT8 quantization, while FP8 methods need Hopper GPUs.
Approach
Use an INT8 data flow in which every linear and non-linear operator reads and writes INT8 tensors with per-block scale factors (B=32), so that quantization is applied on-chip and the weight gradient is also computed in INT8 with per-block scales that keep accuracy under channel-wise outliers.
Key results
- Jetfire matches FP baseline BLEU (26.49 vs 26.49) on Transformer-Base and Top-1 within 0.1% on DeiT-tiny/small, and improves DeiT-base by 0.4%, while per-tensor quantization fails to converge on DeiT.
- On GPT2-base/medium/large, Jetfire reaches validation loss 2.8597/2.4195/2.4696 versus FP 2.9074/2.6612/2.5993, and GLUE within 0.3% of FP.
- End-to-end transformer block speedup of 1.42x overall for hidden size 4096 on RTX 4090 and up to 1.49x activation memory reduction (24 layers, BS=1).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use per-block quantization instead of per-token quantization (ablation) | Jetfire · GPT2 activation quantization | quantization error | Helped |
| Use per-block quantization rather than per-channel quantization (ablation) | Transformer pretraining · Transformer linear layers | quantization error | No clear effect |
| Set Triton block size to 64 × 64 and CUDA block size to 128 × 32 × 128 (ablation) | Jetfire operators · Kernel speed tests | kernel speed | Helped |
| Use INT8 data flow and per-block quantization rather than FP16 training | Jetfire · WMT14 En-De | BLEU: 26.49 -> 26.49 | No clear effect |
| Use Jetfire rather than FP training | DeiT-Base pretraining · ImageNet-1K | Top-1 accuracy: 75.67 -> 76.03 | Helped |
| Use Jetfire rather than FP training | GPT2 pretraining · OpenWebText validation, GPT2-base | validation loss: 2.9074 -> 2.8597 | Helped |
| Use Jetfire rather than FP training | GPT2 fine-tuning · GLUE average score, GPT2-base | GLUE score: $78.50_{0.45}$ -> $78.18_{0.34}$ | Hurt |
| Use Jetfire rather than SwitchBack | GPT2 transformer layer, hidden size 4096 · GPT2 transformer layer | overall speedup: 1.27 -> 1.42 | Helped |
| Use Jetfire rather than SwitchBack | GPT2 transformer layer, hidden size 1024 · GPT2 transformer layer | overall speedup: 1.07 -> 1.07 | No clear effect |
| Use Jetfire rather than SwitchBack | GPT2 activation storage, 24 layers, BS=1 · GPT2 | activation memory reduction: 1.24 -> 1.49 | Helped |
Surprising findings
- Jetfire's validation loss on GPT2 models is lower than the FP16 baseline, which the authors attribute to a regularization effect of quantization.
- Per-tensor quantization fails to converge on DeiT-tiny/small/base and SwitchBack loses up to 4.3% GLUE accuracy on GPT2-large fine-tuning.
Limitations
- End-to-end speedup drops below 1 (0.97 forward, 1.07 overall) at hidden size 1024, so gains depend on matrix size.
- Only hardware evaluated is RTX 4090 and RTX 3090 consumer GPUs; no large-scale multi-GPU pretraining results.
- Multi-head attention is kept in FP16 via FlashAttention, so attention is not fully quantized.
Experimental setup
Datasets: WMT 14 En-De ImageNet1K OpenWebText GLUE
Benchmarks: WMT 14 En-De ImageNet1K GLUE OpenWebText validation loss
Metrics: BLEU Top-1 accuracy validation loss GLUE score kernel speedup end-to-end speedup activation memory reduction quantization MSE
Baselines: FP16 training SwitchBack per-tensor quantization GDA DAQ
Builds on: SwitchBack TransformerEngine FlashAttention Triton
Closest papers by meaning
- The Quantization Benefits of Residual-Free Transformers (NeurIPS 2026 · similarity 0.87)
- A General and Efficient Training for Transformer via Token Expansion (CVPR 2024 · similarity 0.86)
- FrameQuant: Flexible Low-Bit Quantization for Transformers (ICML 2024 · similarity 0.85)
- Accelerating Transformer Pre-training with 2:4 Sparsity (ICML 2024 · similarity 0.85)
- Simplifying Transformer Blocks (ICLR 2024 · similarity 0.85)
- Towards Next-Level Post-Training Quantization of Hyper-Scale Transformers (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.