Single-Stage Huffman Encoder for ML Compression

TL;DR

The paper tackles latency and bandwidth bottlenecks in distributed LLM training/serving by replacing the traditional three-stage Huffman encoder with a single-stage approach that uses fixed codebooks derived from the average symbol distribution observed across prior batches. This leverages observed cross-shard similarity in tensor histograms, particularly FFN activation tensors, to approximate the true distribution without on-the-fly frequency analysis. Empirical results on Gemma 2B show the average PMF closely matches shard-level distributions (low $D_{KL}$), enabling fixed codebooks to achieve compressibility within $0.5\%$ of per-shard Huffman and within $1\%$ of the ideal Shannon bound. The codebooks can be precomputed off the critical path and shared among nodes, yielding a true one-stage encoder that eliminates codebook transmission and reduces latency in die-to-die communication scenarios.

Abstract

Training and serving Large Language Models (LLMs) require partitioning data across multiple accelerators, where collective operations are frequently bottlenecked by network bandwidth. Lossless compression using Huffman codes is an effective way to alleviate the issue, however, its three-stage design requiring on-the-fly frequency analysis, codebook generation and transmission of codebook along with data introduces computational, latency and data overheads which are prohibitive for latency-sensitive scenarios such as die-to-die communication. This paper proposes a single-stage Huffman encoder that eliminates these overheads by using fixed codebooks derived from the average probability distribution of previous data batches. Through our analysis of the Gemma 2B model, we demonstrate that tensors exhibit high statistical similarity across layers and shards. Using this approach we achieve compression within 0.5% of per-shard Huffman coding and within 1% of the ideal Shannon compressibility, enabling efficient on-the-fly compression.

Figures (4)

• Figure 1: Probability Mass Function (PMF) of FFN1 activation.
• Figure 2: Compressibility of FFN1 activation shards using Huffman codes.
• Figure 3: KL divergence of FFN1 activation shards from the average PMF.
• Figure 4: Compressibility using Huffman codes derived from the average distribution.