Warning: very long post. TLDR: this post answers some questions I had about generating text with full, unquantized Falcon-180B under budget constraints.

What is the goal

The goal is to benchmark full, unquantized Falcon-180B. I chose Falcon-180B because it is the biggest open-source model available currently. I also do not use any optimization such as speculative decoding or any kind of quantization, or even torch.compile. I benchmark both for small and large context sizes. I aim for maximum utilization of the available GPUs. I use 3090 cards for all experiments, as they are easy to find in used condition (cost around 700$) and have 24GB of memory.

About the model

The Falcon-180B has 80 transformer layers, the weights are around ~340GB. Its maximum context size is 2048, so whenever I say small context size, I mean around 100 tokens, and whenever I say large context size, I mean 2048 tokens.

Experiment setup

Every LLM can be roughly split into three parts:

1. begin - which converts the tokens into continuous representation (this is usually the embeddings)
2. mid - which is a series of transformer layers. In the case of Falcon-180B we have 80 transformer layers
3. end - which converts the intermediary result into a prediction for the next token (this is usually the LM head)

I converted the Falcon-180B into separate pth file for each of those parts, so for Falcon-180B I have 82 .pth files (one for begin, one for end, and 80 for the transformer layers).

This allows me to save disk space, because for example if a given node is going to run layers 5 to 15, it only needs the weights for those particular layers, there is no need to download several big safetensors files and only read parts of them, instead we aim to store only exactly what is needed for a given node.

I also refactored Falcon-180B so that I can run parts of the model as a normal PyTorch module, e.g. you can run layers 0 to 5 as a normal PyTorch module. This allows me to run it distributed on heterogeneous hardware, e.g. add machines with other cards (which have very little memory) to the computation.

The experiments are being run in distributed mode, with multiple nodes (PCs) having different number of cards, so there is some network overhead, but all nodes are connected to the same switch. In my experiments, I found that the network overhead is about ~25% of the prediction time. This could be improved by using a 10Gbit switch and network cards or Infiniband, but 1Gbit network is the best I could do with the available budget.

Questions

How many layers can you fit on a single 3090 card?

I can load around 5 layers of the Falcon-180B, which take up around 21GB of memory, and the rest 3GB is left for intermediary results. To load all the weights of Falcon-180B on 3090 cards, you would need 16 cards, or 11k USD, assuming used 3090s cost around 700$, although you can also find them for 500$ at some places.

How long does it take to load the state dict of a single node on the GPU?

~3.5s

For 5 layers, it takes ~3.5 seconds to move the state dict from the CPU to the GPU.

How long does it to take to forward a small prompt through a single transformer layer?

~10ms

Since we have 80 layers, the prediction would take at least ~800ms. When you add the begin, end and the data transfer overhead, we go around a little bit more than 1s per token.

How long does it to take to forward a large prompt through a single transformer layer?

~100ms

Since we have 80 layers, the prediction would take at least ~8000ms, or 8 seconds. When you add the begin, end and the data transfer overhead, we go around a little bit more than 10s per token.

How many 3090s do I need to run Falcon-180B with a large prompt?

8

At first glance, it may seem like you need 16 3090s to achieve this, but shockingly, you can do with only 8 3090s and have the same speed of generation!

Why? Because you can reuse the same GPU multiple times! Let me explain what I mean.

Let's say on node0 you load layers 0-5 on the GPU, on node1 you load layers 5-10 on the GPU, etc. and on node7 you load layers 35-40. After node0 does its part of the prediction (which will take ~500ms), it sends to the next node, and while the other nodes are computing, instead of sitting idle, it starts to immediately load layers 40-45 to the GPU, which are pre-loaded in the CPU memory. This load will take around ~3.5 seconds, while the prediction of the other nodes will take ~4s, and since these two processes happen in parallel, there'll be no added time to the total inference time, as each node uses the time in which the other nodes are computing to load future layers to the GPU.

That's insane because in under 6k USD you can 8 3090s and have Falcon-180B running at maximum context size with 10s/token. Add in another 4k USD for the rest of the components, and under 10k USD you can have Falcon-180B running at decent speed.

Implementation details

I separated the project into 4 small libraries with minimal third-party dependencies:

1. One for converting the weights into a separated weights format
2. One for running a node with reloading of future layers
3. One for sampling the results
4. One with Falcon stuff needed to run only parts of it as PyTorch modules. I did regression tests to ensure I have not broken anything and my implementation conforms to the original one

If there is sufficient interest, I may package and open-source the libraries and notebooks.

Future work

I plan to convert other models into the same format and refactor them so that different parts of the model can be used as normal PyTorch modules. Here's which models are currently on my TODO list:

1. Goliath-120b
2. Llama2
3. Mistral
4. Yi

etc.

If the community is interested, I can open-source the whole project and accept requests for new models to be converted into this format.

Thank you for your attention and sorry once again for the long post.