HALO: Hadamard-Assisted Lower-Precision Optimization for LLMs

TL;DR

This work tackles the challenge of quantized training for large language models by introducing HALO, which uses strategically placed Hadamard rotations to mitigate outliers and enable all forward and backward matrix multiplications to run in low precision. The method defines two HALO levels for different dynamic ranges, integrates with HQ-FSDP to reduce communication, and supports both full and parameter-efficient fine-tuning. Empirical results on LLama3-8B and Qwen demonstrate near full-precision accuracy across several tasks, with end-to-end speedups up to about 1.41x on RTX 4090 GPUs, validating HALO as a practical, efficient approach to fully quantized LLM fine-tuning. The work also delivers hardware-conscious design elements, including efficient CUDA kernels and activation-memory reductions, with implications for broader deployment of low-precision LLM training.

Abstract

Quantized training of Large Language Models (LLMs) remains an open challenge, as maintaining accuracy while performing all matrix multiplications in low precision has proven difficult. This is particularly the case when fine-tuning pre-trained models, which can have large weight and activation outlier values that make lower-precision optimization difficult. To address this, we present HALO, a novel quantization-aware training approach for Transformers that enables accurate and efficient low-precision training by combining 1) strategic placement of Hadamard rotations in both forward and backward passes, which mitigate outliers, 2) high-performance kernel support, and 3) FSDP integration for low-precision communication. Our approach ensures that all large matrix multiplications during the forward and backward passes are executed in lower precision. Applied to LLAMA-family models, HALO achieves near-full-precision-equivalent results during fine-tuning on various tasks, while delivering up to 1.41x end-to-end speedup for full fine-tuning on RTX 4090 GPUs. HALO efficiently supports both standard and parameterefficient fine-tuning (PEFT). Our results demonstrate the first practical approach to fully quantized LLM fine-tuning that maintains accuracy in 8-bit precision, while delivering performance benefits. Code is available at https://github.com/IST-DASLab/HALO.

Figures (6)

• Figure 1: Largest magnitudes of first 512 channels in output gradients, or errors, of the mlp output (down_proj) in the 15-th layer of Llama3-8B model over 128 tokens in the 60th step of the ViGGO fine tuning (see Appendix \ref{['appendix:hadamard_affect']} for other data types). The outliers are propagated across columns and can be mitigated after applying left-hand-side Hadamard transformations.
• Figure 2: Cosine similarity between the weight gradients of the baseline model (BF16) and the quantized model when quantization is applied during the (a): backward pass, (b): forward pass, and (c): forward pass with Hadamard transformation in a single fine-tuning step of Llama3-8B on the GSM8K dataset. Compared to the backward pass, the forward pass is more sensitive to quantization. We improve the results by applying Hadamard transformation during the forward pass ($\mathbf{\stackrel{\mathbf{H}}{\mathbf{F}}}$). For each case, we present the weighted average (over # parameters) for all linear modules in each layer.
• Figure 3: Left: Hadamard transformations in HALO levels during the input gradient calculation. Hadamard transformations are already applied to the weights during the forward pass. Right: Forward pass in HQ-FSDP for two GPUs: Each GPU performs a right-hand side Hadamard transformation and computes the absolute maximum (AbsMax) $s_i$ over its local weight shard $W_i$. Then, all GPUs participate in an All-Reduce operation and compute the global maximum absolute value $s$. Each GPU uses this global value to quantize its own weight shard. Finally, an All-Gather operation is performed on the quantized shards $(W_iH)_Q$. All GPUs use the same scales to quantize the weights in the backward pass after applying Hadamard transformation.
• Figure 4: INT8 forward + backward speedups (over BF16) for $512$ sequence length in Llama3-8B across batch sizes (BS). HALO-2 refers to HALO-2($\stackrel{\mathbf{H}}{\mathbf{F}^{\space}}$, $\prescript{\mathbf{H}}{}{\mathbf{E}}^{\mathbf{H}}$, $\mathbf{G}^{\mathbf{H}}$).
• Figure 5: The effect of applying Hadamard transformations on the left and right hand sides of the the activations, errors, and weights.