Optimization story: Bloom inference

This article gives you the behind-the-scenes of how we made an efficient inference server that powers bloom. inference server that powers https://huggingface.co/bigscience/bloom.

We achieved a 5x latency reduction over several weeks (and 50x more throughput). We wanted to share all the struggles and epic wins we went through to achieve such speed improvements.

A lot of different people were involved at many stages so not everything will be covered here. And please bear with us, some of the content might be outdated or flat out wrong because we're still learning how to optimize extremely large models and lots of new hardware features and content keep coming out regularly.

If your favorite flavor of optimizations is not discussed or improperly represented, we're sorry, please share it with us we're more than happy to try out new stuff and correct our mistakes.

Creating BLOOM

This goes without saying but without the large model being accessible in the first place, there would be no real reasons to optimize inference for it. This was an incredible effort led by many different people.

To maximize the GPU during training, several solutions were explored and in the end, Megatron-Deepspeed was chosen to train the end model. This meant that the code as-is wasn't necessarily compatible with the transformers library.

Porting to transformers

Because of the original training code, we set out to do something which we regularly do: port an existing model to transformers. The goal was to extract from the training code the relevant parts and implement it within transformers. This effort was tackled by Younes. This is by no means a small effort as it took almost a month and 200 commits to get there.

There are several things to note that will come back later:

We needed to have smaller models bigscience/bigscience-small-testing and bigscience/bloom-560m. This is extremely important because they are smaller, so everything is faster when working with them.

First inference (PP + Accelerate)

Results