Quantized Training

[msaroufim]

Inspired by a recent back and forth with @gau-nernst we should add some quantized training recipes in AO for small models (600M param range)

Character.ai recently shared that they're working on quantized training https://research.character.ai/optimizing-inference/ where per @stephenroller they train models from scratch in int8 https://x.com/stephenroller/status/1816636257717436779

Historically we've invested more in QAT which @andrewor14 has led which is more of a technique to reduce perplexity when we do an eventual post training quantization.

Quantized training on the other hand actually quantizes the model at training time and so memory savings are observed both for training and inference

So when discussing quantized training there's a few aspects

1. Weights they can be in one: fp16, fp8, int8, int4 and below
2. Activations most likely limited to fp8, fp16
3. Optimizer can be in one of: fp32, fp16, bf16, fp8, int8 and below

And if one were to ship this work, a bad combination can be validated at small scale (~600M parameter range) but a good idea needs to continuously be tested from (8b to 405b range) so each of these will need loss curves

When choosing the starting point, we could either pretrain a model using quantized training or just finetune it and as long as the loss curves match the fp16 baselines then we are good. We'd also need to of course validate that memory savings are there and what the speedups/slowdowns are.

And while we can merge a lot of the dtype conversion in AO and have some toy training loop in AO what I'm more optimistic about is having some end to end trainig recipe in https://github.com/pytorch/torchtitan @awgu and an end to end finetuning recipe https://github.com/pytorch/torchtune @ebsmothers @joecummings

[gau-nernst]

Just want to add, there is also activation/computation dtype and gradient dtype. In my exploration, I still use activation/computation in BF16 and gradient in BF16 to match weight-only quant inference. Activation/computation can be in lower precision dtype also, such as INT8 act - INT8 weight to match dynamic quant inference, or FP8 act - FP8 weight to match current FP8 training recipe.

Lower precision gradient might not be possible? Will need to check existing works on this.

[gau-nernst]

Some extra info for future reference

• Quantized training in JAX: blog post github
• INT8 training by Q-GaLore: paper github

For evaluating the effectiveness of quantized training

• For INT8 quantized training, it's better to compare against BF16 training from scratch / pre-training, since INT8 PTQ is already very good.
• For fine-tuning workflow, compare INT4 quantized training against the current INT4 QAT recipe.

Digging into AQT INT8 (update as I read more). Many things can be customized, but the basic config is:

• Quantize both weight and activation to INT8 to use INT8 matmul. Channel-wise symmetric quantization so that the inner matmul can be INT8, and the scale is absorbed post-matmul.
• Signed number scale to [-127,127]. Unsigned numbers (e.g. after ReLU) scale to [0, 255]. Q: how does UINT8 x INT8 matmul work?
• Calculating scale (https://arxiv.org/abs/2105.03536):
  • For weights: absmax scaling.
  • For activations : train for 20% without activation quantization, keep track of absmax EMA. For the remaining 80%, apply activation quantization with that absmax EMA, and don't update activation scale anymore. (Note: the paper was for ResNet-50. Might not be needed for transformers using LayerNorm/RMSNorm)
• Seems like master weight is still in BF16 -> similar to QAT. Backward is quantized too -> similar to current FP8 recipe.

[gau-nernst]

Found this interesting paper - Jetfire. ICML 2024 poster spotlight. With code release and custom CUDA kernels :open_mouth:

• Paper: https://arxiv.org/abs/2403.12422
• Code: https://github.com/thu-ml/Jetfire-INT8Training

INT8 for everything, including activations and gradients. Tile-wise quantization. Also use 127 for scaling. Master weight in FP32.

Which also led to me an earlier paper - SwitchBack. Timm Dettmers is one of the authors :laughing:.

• Paper: https://arxiv.org/abs/2304.13013
• Code: https://github.com/bitsandbytes-foundation/bitsandbytes/blob/main/bitsandbytes/nn/triton_based_modules.py

Dynamic quantization for everything (weight is still in high precision). Row-wise quant (i.e. batch dim - per token) for activation (forward) and grad output (backward). Tensor-wise quant for weight. INT8 matmul for forward (Y = X @ W.T) and input grad backward (X_grad = Y_grad @ W), while weight grad is FP16 matmul (W_grad = Y_grad.T @ X)

[jerryzh168]

@gau-nernst thanks for the pointers, feels like these are good motivations to enable training with AffineQuantizedTensor since it will be general to support all kinds of quantization (per block, row-wise, per token) and both for dynamic quant and weight only quant. cc @andrewor14