Quad Length Codes for Lossless Compression of e4m3

1. Problem Statement

Training and serving Large Language Models (LLMs) involves massive data parallelization across multiple accelerators (e.g., TPUs). This process relies heavily on collective communication operations (AllReduce, AllGather, etc.), which are frequently bottlenecked by network bandwidth.

While lossless compression can alleviate this bottleneck, existing solutions have significant trade-offs:

• Huffman Codes: Optimal for compression efficiency by exploiting symbol frequency distributions. However, they require deep binary tree traversals for decoding, resulting in slow, bit-sequential decoding and high hardware complexity.
• Universal Codes (e.g., Exponential-Golomb): Faster to decode because the code length is self-delimiting, but they do not exploit specific symbol frequency distributions, leading to suboptimal compression ratios.

The goal is to develop a coding scheme that balances high compression efficiency (like Huffman) with fast, non-sequential decoding and low hardware complexity (like Universal codes).

2. Methodology: Quad Length Codes

The authors propose Quad Length Codes, a hybrid approach designed specifically for the e4m3 (4-bit exponent, 3-bit mantissa) floating-point data type used in ML activations and weights.

Core Mechanism:
Instead of assigning a unique code length to every symbol (as Huffman does), the scheme groups the 256 possible 8-bit symbols into 8 distinct "areas" based on their probability frequency.

• Prefix Bits: The first 3 bits of the code serve as an "Area Code" to identify which of the 8 areas the symbol belongs to.
• Variable Lengths: Each area contains a specific number of symbols and requires a specific number of additional bits to encode the symbol index within that area.
• Resulting Code Lengths: The scheme utilizes only 4 distinct code lengths (e.g., 6, 7, 8, and 11 bits), significantly simplifying the decoder logic compared to Huffman's 13+ distinct lengths.

Adaptability:
The paper demonstrates that the scheme can be adapted to different data distributions:

1. Standard Distribution (FFN1): Uses a fixed partitioning (Table 1) where the first 5 areas have 8 symbols each.
2. Skewed Distribution (FFN2): In cases where one symbol (e.g., zero) is significantly more frequent, the partitioning is adjusted (Table 2) to assign shorter codes (e.g., 4 bits) to the most frequent symbols, improving compression.

Implementation:

• Encoder: Uses a 256-entry Look-Up Table (LUT). Input symbols are sorted by frequency, mapped to indices 0–255, and assigned codes based on the Quad Length scheme.
• Decoder: Reads the first 3 bits to determine the code length, reads the remaining bits, and uses a second 256-entry LUT to map the encoded value back to the original symbol. This avoids tree traversal entirely.

3. Key Contributions

• Novel Coding Scheme: Introduction of Quad Length Codes, which trade a marginal amount of compression efficiency for a drastic reduction in decoding latency and hardware complexity.
• Hardware Optimization: By limiting code lengths to only 4 distinct values and using a flat LUT instead of a deep binary tree, the scheme simplifies hardware implementation, making it suitable for high-throughput ML accelerators.
• Adaptive Partitioning: A methodology to re-configure the "areas" and symbol counts based on the specific Probability Mass Function (PMF) of the data (e.g., handling skewed distributions where zero is dominant).
• Empirical Validation: Extensive analysis on Gemma 2B model tensors (weights, activations, gradients) during Supervised Fine Tuning (SFT).

4. Results

The authors evaluated the scheme on FFN1 and FFN2 activation tensors using the e4m3 data type.

• Trade-off: The standard Quad Length scheme achieves 13.9% compression, which is approximately 2.0% lower than Huffman's 15.9%.
• Benefit: Despite the slight loss in compression, the scheme offers significantly faster decoding and simplified hardware because it eliminates the need for bit-sequential tree traversal.
• Adaptation: When the coding scheme was adapted to the specific skewed distribution of FFN2 activations, the compression improved to 19.0%, narrowing the gap with Huffman (23.2%) while retaining the speed benefits.

5. Significance

This work addresses a critical bottleneck in LLM infrastructure: network bandwidth during collective operations.

• Practicality: While Huffman is theoretically optimal, its decoding latency is often prohibitive in high-speed, parallel hardware environments. Quad Length Codes offer a pragmatic "sweet spot" that is "good enough" in compression but "much better" in speed.
• Scalability: The use of simple LUTs and fixed-length prefixes makes the decoder highly amenable to parallelization and hardware acceleration (ASIC/FPGA), which is essential for the next generation of AI accelerators.
• Future Direction: The authors suggest that the current partitioning is empirical and propose future work to develop a mathematical formulation to optimize the number of areas and symbol distribution for any given dataset.

In summary, Quad Length Codes provide a robust, hardware-friendly alternative to Huffman coding for ML workloads, prioritizing decoding speed and implementation simplicity with a minimal sacrifice in compression efficiency.