Hugging Face BLOOM Inference Optimization

Hugging Face reduced inference latency by 5x and increased throughput by 50x for the BLOOM model through a series of iterative optimizations. The final architecture utilizes PyTorch combined with Tensor Parallelism (TP), a custom CUDA kernel for attention, and torch.jit.script for kernel fusion.

Transitioning from Pipeline to Tensor Parallelism

Initial inference for BLOOM (176B parameters, 352GB in bf16) was implemented using Pipeline Parallelism (PP) via the accelerate library's device_map="auto". In PP, each GPU owns a specific set of layers, processing data sequentially and handing it off to the next GPU.

To reduce latency, Hugging Face transitioned to Tensor Parallelism (TP), where each GPU owns a portion of the weights for every layer, allowing all GPUs to work simultaneously. This shift resulted in a dramatic performance increase:

  • Latency: Dropped from 300ms/token to 91ms/token.
  • Throughput: Increased to 10 requests per second (RPS).

While TP introduces communication overhead via ncclAllReduce, the ability to process batches (where batch size 1 and 32 often have similar latency) significantly improved overall throughput.

PyTorch-Based Optimizations

Beyond parallelism strategies, several low-level PyTorch optimizations were implemented to remove bottlenecks identified through profiling with TensorBoard.

Kernel Fusion with torch.jit.script

The Gelu operator originally launched multiple element-wise kernels, causing excessive tensor copies. By applying @torch.jit.script to the bloom_gelu_forward function, Hugging Face fused these into a single kernel operation, reducing latency from 91ms/token to 81ms/token.

Efficient PyTorch Implementation

  • ALiBi Optimization: Position embeddings were previously calculated in too many locations. Centralizing this calculation resulted in a 10x speedup for that specific operation.
  • Reducing Tensor Copies: Profiling revealed that the attention path was heavily burdened by reshape and transpose operations. Reworking the weights and the KV cache (the "past") removed these unnecessary copies.

Custom CUDA Kernels and Hardware Acceleration

To further optimize the hot path where torch.jit.script was insufficient, Hugging Face developed a custom CUDA kernel to fuse the masked fill and softmax operations.

Specifically, the kernel optimizes the following sequence:

  1. masked_fill_ of attention scores using the attention mask.
  2. softmax calculation on float32 for stability.

By limiting upcasting to only the necessary sums and accumulations within the kernel, latency was further reduced from 81ms/token to 71ms/token.

Webserver Architecture and Request Handling

To serve a diverse set of user requests with varying parameters and lengths, Hugging Face implemented a flexible batching system:

  • Inter-process Communication: Because torch.distributed requires separate processes, the server uses Redis pub/sub to distribute raw strings to all processes.
  • Custom Generation Loop: The standard generate function was replaced with a custom implementation that applies different parameters (e.g., sampling, top-p) to each member of a batch.
  • Dynamic Batch Extraction: To prevent short requests from being delayed by long requests in the same batch, the server extracts and returns completed requests as soon as they reach their token limit, rather than waiting for the entire batch to finish.

Evaluated but Discarded Approaches

Throughout the optimization process, several other routes were explored:

  • JAX/Flax on TPUs: While parallelism was easier to implement, the team encountered significant stability issues with Ray and TPU worker communication, and lacked fine-grained control over compilation.
  • DeepSpeed: Provided impressive results similar to the final iteration but suffered from stability issues, including regular kernel crashes (CUDA illegal access) under stress.
  • Rust Implementation: A version was written in Rust using tch-rs for better concurrency control. However, it was discovered that the perceived performance gain was actually due to a profiler being left active in the PyTorch benchmarks.
  • ONNX/TensorRT: These were deemed too rigid for the required flexibility of the text-generation loop and the need to keep tensors on the GPU for logits computation.

Sources