SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexander Borzunov*  
HSE University, Yandex
Max Ryabinin  
HSE University, Yandex
Artem Chumachenko
Neiro.ai
Dmitry Baranchuk  
Yandex
Tim Dettmers  
University of Washington
Younes Belkada Hugging Face
Pavel Samygin  
Yandex School of Data Analysis
Colin Raffel Hugging Face
Large language models (LLMs) are useful in many NLP tasks and become more capable with size, with the best open-source models having over 50 billion parameters. However, using these  $50\mathrm{B}+$  models requires high-end hardware, making them inaccessible to most researchers. In this work, we investigate methods for cost-efficient inference and fine-tuning of LLMs, comparing local and distributed strategies. We observe that a large enough model  $(50\mathrm{B}+)$  can run efficiently even on geodistributed devices in a consumer-grade network. This could allow running LLM efficiently by pooling together idle compute resources of multiple research groups and volunteers. We address two open problems: (1) how to perform inference and fine-tuning reliably if any device can disconnect abruptly and (2) how to partition LLMs between devices with uneven hardware, joining and leaving at will. In order to do that, we develop special fault-tolerant inference algorithms and load-balancing protocols that automatically assign devices to maximize the total system throughput. We showcase these algorithms in PETALS<sup>1</sup> — a decentralized system that runs Llama 2 (70B) and BLOOM (176B) over the Internet up to  $10\times$  faster than offloading for interactive generation. We evaluate the performance of our system in simulated conditions and a real-world setup spanning two continents.
In recent years, the NLP community has found that pretrained language models greatly accelerated progress on many research problems through either fine-tuning (Radford et al., 2018) or simple prompting (Brown et al., 2020). Their quality tends to improve as we increase model scale (Radford et al., 2019; Kaplan et al., 2020). Following this trend, modern language models often have hundreds of billions of parameters (Brown et al., 2020; Rae et al., 2021; Zeng et al., 2021; Kim et al., 2021).
Most recently, several research groups open-sourced their pretrained LLMs with over 50B parameters (Zhang et al., 2022; BigScience, 2022a; Touvron et al., 2023a,b). However, they are still difficult to use due to the sheer size in terms of parameters. For example, OPT-175B and BLOOM-176B need over 350 GB accelerator memory for inference and even more for fine-tuning. As a result, even basic inference for these LLMs requires multiple high-end GPUs or multi-node clusters. Recent studies
*Correspondence to: borzunov alexander@gmail.com
<sup>1</sup>PETALS source code and documentation are available at https://petals.dev
Figure 1: A high-level overview of our system design. Servers store pretrained LLM layers and temporarily hold attention caches for inferencing. Clients hold embedding layers and learned prompts/adapters (if used). Arrows denote temporary chains formed for inference.
propose algorithms for running large models with more affordable hardware (Pudipeddi et al., 2020; Ren et al., 2021), e.g. by offloading parameters to RAM. However, as we show in Section 3.1, these techniques are inefficient in many use cases, such as LLM-based chatbots and search engines.
In this work, we search for a more cost-effective way of running pretrained LLMs in their main use cases: inference, in-context learning, and fine-tuning. We analyze latency and throughput for these use cases and determine which factors become dominant for very large models. Notably, for models with over 50B parameters, communicating activations over a slow network can be faster than swapping layers from local RAM or SSD. Based on these observations, it should be possible to run LLMs cost-effectively by pooling together commodity hardware over the Internet.
However, existing LM algorithms are not designed to run inference with unreliable devices or high-latency networks. To bridge this gap, we formulate a novel algorithm for fault-tolerant distributed autoregressive inference of very large models. Using dual attention caches, this algorithm can quickly recover from a failed server and reassign the load to one or more replacement servers. Finally, to make sure that there are enough servers for every part of the model, we develop a decentralized load-balancing algorithm that assigns transformer blocks to every server to maximize the total system throughput. The fully decentralized nature of these protocols allows participants to add or remove their devices at any point, making optimal use of GPU idle time.
We summarize the main contributions of this work as such:
- We analyze the problem of cost-efficient LLM inference and propose a novel algorithm that can inference large  $(50\mathrm{B}+)$  language models on distributed unreliable devices. To the best of our knowledge, this is the first algorithm that can inference LLMs with  $50\mathrm{B}+$  parameters in this setup.
- Using this algorithm, we develop PETALS — a decentralized system for inferencing and fine-tuning LLMs over the Internet. The system allows users to run inference and fine-tuning over a swarm of unreliable devices with the same correctness guarantees as when running locally. The system runs persistently with the help of volunteers.
- We benchmark the performance of the proposed algorithms on Llama 2 (70B) (Touvron et al., 2023b) and BLOOM (176B) (BigScience, 2022a). We run experiments in controlled conditions, with simulated network latency and server failures, and in the actual geo-distributed system spanning two continents. With realistic network speeds, our distributed algorithms perform autoregressive generation  $\geq 10\times$  faster than local offloading.
2 Background: efficient training and inference
There is a wide variety of methods optimizing training and inference for most deep learning workloads. Here, we focus on two areas relevant for our analysis: model parallelism and parameter offloading.
2.1 Model parallelism
Model parallelism is a family of distributed training algorithms that assigns each device to hold a subset of model parameters, run a subset of computations and communicate output activations. Tensor parallelism assigns each device to compute a subset of each model layer (e.g., a subset of neurons), then communicate results between each other and proceed to the next layer (Krizhevsky et al., 2012;
Ben-Nun & Hoeffer, 2019; Tang et al., 2020). Each device performs a symmetric computation, applied to a different slice of model weights, which makes tensor parallelism compatible with MPI-based communication. In turn, the main performance overhead of this strategy comes from all-to-all communication (and synchronization) after each layer (Krizhevsky, 2014).
Pipeline parallelism reduces the communication overhead by assigning each device with one or several full layers (Huang et al., 2019; Narayanan et al., 2019; Yang et al., 2019). During the forward pass, each stage applies its subset of layers to the inputs supplied by the previous stage, then sends the outputs of the last layer to the next stage. For the backward pass, this process is reversed, with each pipeline stage passing the gradients to the same device that previously supplied it with input activations. To better utilize the available devices, the pipeline must process multiple microbatches per step, allowing each stage to run in parallel on a different batch of inputs. Even with optimal execution, some of the pipeline stages will remain idle some of the time (Huang et al., 2019).
Both of these strategies are actively used for training LLMs. Real-world distributed training systems usually combine multiple forms of parallelism depending on hardware and network type (Narayanan et al., 2021; Rajbhandari et al., 2020; Jia et al., 2019). Tensor parallelism is typically used within a single multi-GPU server or closely interconnected TPU cores (Narayanan et al., 2021; Shazeer et al., 2018). In turn, pipeline parallelism is used to connect multiple servers (Narayanan et al., 2021). Recent works demonstrate that model parallelism can be used for cost-efficient pre-training of LLMs by pooling together idle GPU devices (Athlur et al., 2022; Wang et al., 2022; Kuszmaul, 2022; Yuan et al., 2022; Ryabinin et al., 2021).
2.2 Offloading
Parameter offloading relegates model parameters from accelerator memory to a slower but cheaper storage: typically RAM or SSD (Pudipeddi et al., 2020; Ren et al., 2021; Rajbhandari et al., 2021). When using the model, parameters are loaded to the accelerator just-in-time for computation, one or few layers at a time. In principle, this method allows running large models with a single low-end accelerator as long as there is enough RAM (or SSD) to store the model.
The main drawback of this strategy is having to load and unload through all model parameters for each forward and backward pass, which can be time-consuming. This extra time can be amortized in workloads where model can do a lot of useful computations for each time a parameter is loaded. In practice, using offloading to run a single token through the OPT-175B on one GPU in the best-case scenario of hardware and bandwidth² would require 11 seconds per forward pass, or twice that for training. As we show in Section 4, real-world performance is significantly slower.
Pudipeddi et al. (2020) circumvents this by training with very large batches, and hence, increasing the computation. In turn, Ren et al. (2021); Rajbhandari et al. (2021) reduce the overhead by overlapping communication and computation, that is, doing useful computation for the current layer while waiting for the transfer of the next layer to finish. Some of these systems Ren et al. (2021) also partition offloaded parameters between devices. However, unlike model-parallel training, distributed offloading still requires each device to compute the full model.
Using pretrained large language models for NLP tasks consists of two main workloads: inference and fine-tuning. The inference workload typically consists of encoding an input text, then generating tokens autoregressively. In turn, fine-tuning requires updating either all of the model's parameters or (more commonly for large models) a small set of trainable weights (e.g., adapters or soft prompts) by backpropagation. These two workloads also cover more advanced use cases:
- Manually engineering prompts for a given task, then deploying the model with these prompts.
- Fine-tuning with adapters (Hu et al., 2021; Houlsby et al., 2019; Liu et al., 2022b) or "soft" prompts (Liu et al., 2021b; Lester et al., 2021; Liu et al., 2021a) and inferencing fine-tuned models.
- Distillation into a smaller task-specific model for faster inference (Schick & Schütze, 2021).
${}^{2}$  Specifically,16-bit parameters,PCIe gen. 4 at  ${31.5}\mathrm{{GB}}/\mathrm{s}$  (16 lanes),infinite compute and memory bandwidth.
Counter-intuitively, we found that inference is more challenging than fine-tuning for cost-efficient setups. To that end, we dedicate most of this section to inference-specific problems. As for fine-tuning, we describe a way to support arbitrary parameter-efficient fine-tuning in Section 3.4.
3.1 Performance bottlenecks of LLM inference
Unlike training, autoregressive LLM inference cannot be done with a single pass through the model. Instead, the model needs to process one token at a time, pass it through the entire model, then generate the next token and repeat the process. In case of model parallelism, training an  $n$ -layer $^3$  model on a sequence of  $t$  tokens needs  $O(n)$  communication rounds, while generating the same sequence needs  $O(n \cdot t)$  rounds, making it more susceptible to network latency. Similarly with parameter offloading, generating a sequence of  $t$  tokens needs loading every layer  $t$  times, which also takes  $O(n \cdot t)$  time.
The other problem of autoregressive generation is dealing with attention for past tokens (Vaswani et al., 2017). During an inference step  $t$ , each layer needs to attend to  $t - 1$  previous attention keys and values. Existing inference algorithms store past entries in accelerator memory. Caching half-precision activations of a 2048-token sequence for large models like GPT-3 (Brown et al., 2020) or OPT-175B (Zhang et al., 2022) (with 96 layers of 12288 units each) takes up 9.6 GB GPU memory for each sequence. Offloading these cached values faces the same problems as offloading in general.
An alternative solution is to recompute all previous tokens on every inference step, storing only one set of keys & values at a time. Naturally, this approach needs increasingly more computation with sequence length  $t$ , for a total of  $O(t^{3})$  time for transformer-based models<sup>4</sup>. Surprisingly, this approach is often more efficient than offloaded caching, especially for shorter sequences due to the overhead from loading and storing cache from RAM or SSD.
Parameter offloading can still be efficient when generating large amounts of short sequences in bulk. Each individual sequence still takes a long time to generate, but the system maintains high throughput by running many samples in parallel. Unfortunately, this scenario does not cover many important LLM use cases. For instance, it is incompatible with in-context learning or prompt engineering, where the model needs to process long sequences of training examples (Brown et al., 2020). More importantly, it does not support "interactive" applications where LLM needs to quickly respond to a user input. This rules out many LLM applications such as conversation systems or input completion (e.g. ChatGPT or SmartCompose).
Hence, we explore a new solution based on pipeline-parallelism. A related line of work (Aminabadi et al., 2022) investigates model parallelism to inference LLMs in GPU clusters. However, their approach does not apply to our more affordable setups: cheap "preemptible" instances or connecting existing resources over the Internet. To operate in these conditions, an inference algorithm needs to deal with node preemption, network errors, and high latency.
3.2 Distributed generation with fault tolerance
In this section, we formulate an algorithm for inferencing LLMs in a fleet of unreliable geographically distributed devices connected over the Internet. Each device can act as a server, a client, or both. A client is a node operated by the user, which runs inference or fine-tuning jobs through the swarm of servers. A client only holds input and output embeddings ( $< 3\%$  of model weights for BLOOM-176B) and delegates running transformer blocks (the most expensive computations) to remote servers. A server is a GPU-enabled node holding a set of consecutive transformer blocks and processing requests coming from client nodes.
For simplicity, we assume that every block is hosted on several servers and examine this assumption in the next section. Following this notation, a fault-tolerant algorithm should allow each client to complete an inference job with reproducible results even if some remote servers fail during inference.
As we discuss in Section 3.1, autoregressive generation requires many sequential communication rounds, making it sensitive to network latency. However, if every device stores its past attention
<sup>3</sup>Here and below, the term model layer (or block) refers to one transformer block that typically combines self-attention, a feed-forward network, normalization layers, and a residual connection (Vaswani et al., 2017).
cache, every round only transfers activations for a single token, i.e. several kilobytes of data<sup>5</sup>. We use this model to directly minimize the inference time over possible pipeline configurations. As we show later in Section 4.2, this allows efficient inference over a low-bandwidth Internet connection.
A more challenging problem is how to recover from node and network failures. If a remote server shuts down, any cached attention keys stored on that server will be lost with it. There are two naive solutions to this problem: restarting inference from scratch or recomputing past embeddings on every step. Restarting might be enough at a small scale. However, running  $50\mathrm{B}+$  models may involve many unreliable devices, making it unlikely to generate long sequence without at least one failure. In turn recomputing past attention caches requires communicating past tokens on every communication round, resulting in  $O(n \cdot t^2)$  total data transferred, where  $n$  is the number of pipeline layers and  $t$  is the sequence length. In other words, both these solutions struggle to generate long sequences.
We address this problem by maintaining two types of cache: server-side cache holds past attention keys and values for their layers, like in existing inference algorithms, while client-side cache holds past inputs sent to a given pipeline stage<sup>6</sup>. If a server disconnects, a client can find another server with that pipeline stage and use client-side cache to restore the server state.
The resulting procedure is described in Algorithm 1. For every pipeline stage, the client maintains a heap (priority queue) of servers that hold this stage (and may hold additional stages). The servers in queue are ordered by the network latency, measured from past communication. These queues are maintained through the lifetime of a client. To begin generation, the client runs a beam-search-like procedure to find a sequence of servers that results in the least total inference time under our performance model. When running inference steps, a client keeps track of intermediate activations sent between pipeline stages. If a remote server fails or leaves, the client retrieves the next best server (or multiple servers) and requests it to restore the attention state from the client's cached activations.
When servers fail, the algorithm needs to send  $O(t)$  data (in one round) for each failed server and compute only the stages held by the failed servers. This can be seen as an interpolation between naive and cached inference, depending on the server failure rate. If none of the servers fail, we recover  $O(n \cdot t)$  communication, similarly to Aminabadi et al. (2022). In turn, if all servers fail after one step, the algorithm effectively performs non-caching generation, which is the best option in that scenario.
In the basic formulation, all communication between pipeline stages is routed through the client, i.e. the client receives the outputs of every pipeline stage, caches it and sends it to the subsequent stage. In practice, it is more efficient to let pipeline stages communicate directly: once the server obtains output activations, it sends them to both client and the subsequent stage. This reduces the total step time since both messages are a few kilobytes in size and can be sent in parallel. To verify that both client and the next pipeline stage received the same set of activations, they can verify the checksums (i.e. hash values) of the received activations asynchronously, without blocking computation.
Algorithm 1 can support greedy inference or any sampling variants (including Holtzman et al. (2020)). However, it requires one more step to support search-based algorithms such as beam search: cache reordering. This allows a client to generate multiple continuations of the same input prefix by cloning its attention cache and dropping less likely hypotheses. We describe beam search in Appendix C.
Shortest path routing. In the Algorithm 1, the find_best_chain function (line 4) selects a sequence of servers that can run the required layers in the least amount of time. To estimate this time we add up two factors: computation time, determined by server's compute throughput ("GPU speed") and the network latency between the client and that server. Servers measure their own compute throughput and share this information with the clients. In turn, clients measure the network latency between them and a given server by "pinging" the candidate servers during routing. If a server runs multiple consecutive blocks, we multiply the computation time by the number of blocks.
To find the best chain of servers, clients find the shortest path between the first and last block, using a graph where edge weights correspond to server inference time, as described in the previous paragraph. To minimize overhead, we do not run pathfinding from scratch on each call to find_best_chain. Instead, clients run lifelong pathfinding in the background and reuse it between inference calls. More specifically, we use the  $\mathbf{D}^*$  Lite (Koenig & Likhachev, 2005) algorithm because it allows clients to quickly adjust paths after a server is banned or leaves the network.
For GPT-3 and OPT-175B, one 12288-dimensional token embedding in 16-bit precision takes up 24 KiB. Here, a pipeline stage is a set of consecutive model layers hosted on one server (as in pipeline parallelism).
Algorithm 1 Generating sequence, client-side code
Input: prefix_tokens, embeddings, known_servers  
1: generated_sequence = list()  
2: cache = dictionary()  
3: streams = dictionary()  
4: chain = find_best_chain(known_servers)  
5: for server ∈ chain do  
6: streams[server] = rpc_inference(server)  
7: cache[server] = list()  
8: end for  
9:  
10: inputs = embeddings(prefix_tokens)  
11: while should_CONTINUE(generated_sequence) do  
12: tail_servers = copy(chain)  
13: while not empty(tail_servers) do  
14: server = tail_servers.pop_left()  
15: try:  
16: ▷ Attempt normal inference  
17: outputs = streams[server].send(inputs)  
18: cache[server].append(inputs)  
19: inputs = outputs  
20: catch ServerFailed:  
21: ▷ Replace the failed server  
22: streams.pop(server).close()  
23: past_input = cache.pop(server)  
24: new_servers = replace_failed_server(  
25: server, past_input, cache,  
26: streams, known_servers)  
27: chain.replace(server, new_servers)  
28: tail_servers.push_left(new_servers)  
29: end while  
30:  
31: logits = compute_logits(outputs, embeddings)  
32: next_token = choose_next(logits) {e.g. greedy}  
33: generated_sequence.append(next_token)  
34: inputs = embeddings(next_token)  
35: end while  
36:  
37: for server ∈ chain do  
38: streams[server].close()  
39: end for  
40: return generated_sequence
Algorithm 2rpc_inference(server)
Input: local_layers, stream  
1: cache = dictionary()  
2: for layer  $\in$  local_layers do  
3: cache(layer) = make_empty()  
4: end for  
5: while not stream Closed() do  
6: inputs = stream.receive()  
7: for layer  $\in$  local_layers do  
8: past_kv = cache(layer]  
9: inputs, new_kv = forward(  
10: layer, inputs, past_kv)  
11: cache(layer].append(new_kv)  
12: end for  
13: stream.send(input)  
14: end while  
Algorithm 3 replace_FAILED_server(...)  
Input: server, inputs, cache, streams, known_servers  
1: known_servers.ban(server)  
2: missing_layers = get_layers(server)  
3: chains = select_by_layer(  
4: known_servers, missing_layers)  
5: chain = find_best_chain(chains)  
6: replacements = list()  
7: while not empty(chain) do  
8: s = chain.pop_left()  
9: try:  
10: streams[s] = rpc_inference(s)  
11: outputs = streams[s].send(input)  
12: replacements.append(s)  
13: cache[s] = inputs  
14: missing_layers.pop(get_layers(s))  
15: inputs = outputs  
16: catch FailedRPC:  
17: known_servers.ban(s)  
18: chains = select_by_layer(  
19: chains, missing_layers)  
20: chain = find_best_chain(chains)  
21: end while  
22: return chain
3.3 Automatic load balancing
In order to run inference or fine-tuning, each server needs to be assigned to a pipeline stage, then reassigned if other servers join or leave the network. For example, if we deploy an LLM on idle compute resources from several data centers or labs, the number of participants may change over time based on the demand. Moreover, servers may have different compute throughput, network bandwidth, and geographical location. To operate in these conditions efficiently, servers should automatically choose which model layers they should serve in a given situation.