BLOOMZ Inference on Habana Gaudi2 Accelerator
TL;DR
Hugging Face has demonstrated that the Habana Gaudi2 accelerator provides faster inference for large language models like BLOOMZ compared to currently available GPUs. For the 176-billion-parameter BLOOMZ model, Gaudi2 is 1.42x faster than the Nvidia A100 80GB.
BLOOMZ Model Capabilities
BLOOMZ is a fine-tuned version of BLOOM, a 176-billion-parameter autoregressive model designed to complete text sequences across 46 languages and 13 programming languages. BLOOMZ specifically improves upon the original BLOOM architecture by enhancing generalization and zero-shot capabilities—the ability to complete tasks on unseen input data without prior training examples.
Deploying BLOOMZ is computationally demanding; in 16-bit precision, a single instance requires 352 GB of memory, necessitating high-performance hardware and optimization libraries to achieve low latency.
Habana Gaudi2 Hardware and Software Stack
Gaudi2 is a second-generation AI hardware accelerator from Habana Labs. Its architecture is designed to perform General Matrix Multiplication (GeMM) and other operations in parallel, which optimizes deep learning workflows for both training and inference.
Hardware Specifications
A single Gaudi2 server contains eight Habana Processing Units (HPUs), each equipped with 96GB of memory, providing the necessary capacity to host very large models.
Software Integration
- SynapseAI™ SDK: Supports PyTorch and DeepSpeed. It includes a graph compiler that optimizes execution through operator fusion, data layout management, parallelization, pipelining, and memory management.
- HPU Graphs: Recently introduced in SynapseAI to support latency-sensitive applications.
- Optimum Habana: A Hugging Face library that acts as a bridge between the Gaudi2 hardware and the Transformers library, simplifying the deployment process.
Inference Benchmarks: Gaudi2 vs. A100
To handle the memory requirements of BLOOMZ, Hugging Face utilized DeepSpeed-inference (via Habana's DeepSpeed fork) to implement model and pipeline parallelism across eight devices.
Latency Results
Benchmarks were conducted using 16-bit precision, greedy generation of 100 tokens, and a key-value cache. The results show Gaudi2 consistently outperforms the Nvidia A100 80GB in latency:
| Model | Number of devices | Gaudi2 latency (seconds) | A100-80GB latency (seconds) | First-gen Gaudi latency (seconds) |
|---|---|---|---|---|
| BLOOMZ (176B) | 8 | 3.103 | 4.402 | / |
| BLOOMZ-7B | 8 | 0.734 | 2.417 | 3.321 |
| BLOOMZ-7B | 1 | 0.772 | 2.119 | 2.387 |
Key Findings:
- 176B Model: Gaudi2 is 1.42x faster than the A100 80GB.
- 7B Model: Gaudi2 is 2.89x faster than the A100 80GB.
- Model Parallelism: Gaudi2 benefits significantly from model parallelism, whereas the A100 is faster on a single device for the smaller 7B model.
Price-Performance of First-Gen Gaudi
For the BLOOMZ-7B model, first-generation Gaudi provides a superior price-performance ratio compared to the A100. While the A100 costs over $30 per hour, first-gen Gaudi (DL1 instance on AWS) costs approximately $13 per hour, while maintaining competitive latency (2.387 seconds).
Implementation and Reproduction
Inference can be performed on complete datasets using a provided script in the optimum-habana repository. The benchmark environment utilized Transformers v4.28.1, SynapseAI v1.9.0, and Optimum Habana v1.5.0.
To reproduce the results, users must install the latest SynapseAI and Gaudi drivers, then install the optimum-habana library and the Habana-specific fork of DeepSpeed (v1.9.0). The execution command utilizes gaudi_spawn.py with flags for DeepSpeed, HPU graphs, and KV cache to optimize generation.