Incredibly Fast BLOOM Inference with DeepSpeed and Accelerate

This article shows how to get an incredibly fast per token throughput when generating with the 176B parameter BLOOM model.

As the model needs 352GB in bf16 (bfloat16) weights (176*2), the most efficient set-up is 8x80GB A100 GPUs. Also 2x8x40GB A100s or 2x8x48GB A6000 can be used. The main reason for using these GPUs is that at the time of this writing they provide the largest GPU memory, but other GPUs can be used as well. For example, 24x32GB V100s can be used.

Using a single node will typically deliver a fastest throughput since most of the time intra-node GPU linking hardware is faster than inter-node one, but it's not always the case.

If you don't have that much hardware, it's still possible to run BLOOM inference on smaller GPUs, by using CPU or NVMe offload, but of course, the generation time will be much slower.

We are also going to cover the 8bit quantized solutions, which require half the GPU memory at the cost of slightly slower throughput. We will discuss BitsAndBytes and Deepspeed-Inference libraries there.

Benchmarks

Without any further delay let's show some numbers.

For the sake of consistency, unless stated differently, the benchmarks in this article were all done on the same 8x80GB A100 node w/ 512GB of CPU memory on Jean Zay HPC. The JeanZay HPC users enjoy a very fast IO of about 3GB/s read speed (GPFS). This is important for checkpoint loading time. A slow disc will result in slow loading time. Especially since we are concurrently doing IO in multiple processes.

All benchmarks are doing greedy generation of 100 token outputs:

Generate args {'max_length': 100, 'do_sample': False}

The input prompt is comprised of just a few tokens. The previous token caching is on as well, as it'd be quite slow to recalculate them all the time.

First, let's have a quick look at how long did it take to get ready to generate - i.e. how long did it take to load and prepare the model:

Here is the throughput in msecs on 8x80GB GPUs:

Throughput in msecs 4x80GB A100: