# SLoRA: Federated Parameter Efficient Fine-Tuning of Language Models

The paper on S-LoRA focuses on scalable serving of <mark style="color:yellow;">Low-Rank Adaptation (LoRA)</mark> adapters for  language models (LMs).&#x20;

{% embed url="<https://arxiv.org/abs/2308.06522>" %}
SLoRA: Federated Parameter Efficient Fine-Tuning of Language Models
{% endembed %}

### <mark style="color:purple;">**Background and Motivation**</mark>

* <mark style="color:blue;">**Pretrain-then-Finetune Paradigm**</mark><mark style="color:blue;">:</mark> LMs are commonly fine-tuned for specific tasks, leading to numerous fine-tuned variants of a single base model.
* <mark style="color:blue;">**LoRA for Fine-Tuning**</mark><mark style="color:blue;">:</mark> LoRA is a parameter-efficient fine-tuning method that updates only low-rank additive matrices (adapter weights), achieving performance comparable to full-weight fine-tuning.
* <mark style="color:blue;">**Challenges in Serving LoRA Adapters**</mark><mark style="color:blue;">:</mark> While fine-tuning has been extensively researched, the efficient serving of these fine-tuned variants, especially at scale, remains underexplored.

### <mark style="color:purple;">**S-LoRA System**</mark>

* <mark style="color:blue;">**Goal**</mark><mark style="color:blue;">:</mark> To scalably serve thousands of LoRA adapters on a single machine or across multiple GPUs.
* <mark style="color:blue;">**Unified Paging**</mark><mark style="color:blue;">:</mark> S-LoRA introduces Unified Paging to manage dynamic adapter weights with varying ranks and KV cache tensors with different sequence lengths in a unified memory pool.
* <mark style="color:blue;">**Tensor Parallelism and CUDA Kernels**</mark><mark style="color:blue;">:</mark> It employs a tensor parallelism strategy and custom CUDA kernels for heterogeneous batching of LoRA computation.
* <mark style="color:blue;">**Efficiency**</mark>: S-LoRA aims to serve multiple adapters concurrently with high throughput and low latency, significantly improving over traditional methods.

### <mark style="color:purple;">**Technical Innovations**</mark>

* <mark style="color:blue;">**Adapter and Base Model Separation**</mark><mark style="color:blue;">:</mark> S-LoRA separates the batchable base model computation from individual LoRA computations to achieve high-throughput multi-adapter serving.
* <mark style="color:blue;">**Memory Management**</mark><mark style="color:blue;">:</mark> The system stores all adapters in the main memory and fetches the required adapters to the GPU memory as needed, efficiently using GPU memory and reducing fragmentation.
* <mark style="color:blue;">**Serving Throughput and Latency**</mark><mark style="color:blue;">:</mark> Compared to existing libraries (like HuggingFace PEFT and vLLM), S-LoRA can increase throughput by up to 4 times and serve significantly more adapters.

### <mark style="color:purple;">**Potential Impact**</mark>

* <mark style="color:blue;">**Scalable Serving of Fine-Tuned Models**</mark><mark style="color:blue;">:</mark> S-LoRA enables scalable serving of many task-specific fine-tuned models.
* <mark style="color:blue;">**Large-Scale Customized Fine-Tuning Services**</mark><mark style="color:blue;">:</mark> It offers the potential for large-scale customized fine-tuning services, which could be a substantial advancement for deploying personalized AI solutions.

### <mark style="color:purple;">**Challenges in Batched Inference**</mark>

* **Memory Management**: Efficiently managing GPU memory is critical due to the simultaneous serving of multiple LoRA adapters. Storing adapter weights outside the GPU and fetching them dynamically presents challenges in memory fragmentation and I/O overhead.
* **Batching Complexity**: The separated computation of many adapters, each with distinct ranks and stored in non-contiguous memory, complicates the batching process.
* **Multiple GPU Usage**: Utilizing multiple GPUs requires novel parallelism strategies to efficiently handle the added LoRA weights and computations, minimizing communication and memory overheads.

### <mark style="color:purple;">**S-LoRA's Contributions**</mark>

* **Unified Paging**: S-LoRA introduces a unified memory pool to manage dynamic adapter weights and KV cache tensors. This approach reduces memory fragmentation and increases batch size.
* **Heterogeneous Batching**: Custom CUDA kernels are employed for efficiently batching different adapters of varying ranks. These kernels operate on non-contiguous memory, aligning with the memory pool design for efficient batched inference.
* **S-LoRA TP (Tensor Parallelism)**: A novel tensor parallelism strategy ensures effective parallelization across multiple GPUs. It achieves minimal communication cost by scheduling communications on small intermediate tensors and fusing large tensors with the communications of the base model.

### <mark style="color:purple;">**Performance Evaluation**</mark>

* **Model Testing**: S-LoRA was evaluated on models ranging from Llama-7B to Llama-70B.
* **Throughput Enhancement**: Compared to HuggingFace PEFT, a leading parameter-efficient fine-tuning library, S-LoRA can enhance throughput by up to 30 times.
* **Comparison with vLLM**: Against the vLLM system, which naively supports LoRA serving, S-LoRA improves throughput by up to 4 times and significantly increases the number of served adapters.

### <mark style="color:purple;">**Impact and Potential**</mark>

* **Scalability**: S-LoRA's ability to serve thousands of adapters on a single GPU or across multiple GPUs with minimal overhead represents a significant advancement in scalable serving of fine-tuned models.
* **Efficiency in Serving Multiple Tasks**: By overcoming the challenges of memory management, batching complexity, and effective use of multiple GPUs, S-LoRA enables the efficient and scalable serving of a wide range of task-specific fine-tuned models.

In conclusion, the S-LoRA paper presents a comprehensive solution for the scalable serving of LoRA adapters in large language models. By addressing key challenges in memory management, batching, and parallelism, S-LoRA notably enhances throughput and scalability, paving the way for the efficient deployment of highly personalized and diverse AI services.

### <mark style="color:purple;">**LM Architecture and Scale**</mark>

* **Architecture**: Most LMs are based on the transformer architecture.
* **Size**: LMs can have several billion to trillion parameters, resulting in disk sizes from gigabytes to terabytes.
* **Computational and Memory Demands**: The sheer scale of LMs leads to significant computational and memory requirements during serving.

### <mark style="color:purple;">**Inference Process in LMs**</mark>

* **Autoregressive Decoding**: Inference requires iterative autoregressive decoding, where the model <mark style="color:yellow;">encodes the prompt in a forward pass and then decodes the output one token at a time</mark>.
* **Sequential Nature**: The <mark style="color:yellow;">sequential nature of decoding each token, attending to the hidden states of all preceding tokens, makes the process slow</mark>.
* **KV Cache**: There's a need to <mark style="color:yellow;">store the hidden states of all previous tokens</mark>, referred to as the KV cache, which adds substantial memory overhead.

### <mark style="color:purple;">**Challenges in Online Serving**</mark>

* <mark style="color:blue;">**Dynamic Requests**</mark><mark style="color:blue;">:</mark> In online settings, handling requests of varying sequence lengths that arrive dynamically is a challenge.
* <mark style="color:blue;">**Fine-Grained Scheduling**</mark><mark style="color:blue;">:</mark> Orca introduces iteration-level scheduling, batching at the token level instead of the request level, allowing for the addition of new requests to the current batch, thus enhancing throughput.
* <mark style="color:blue;">**vLLM's Memory Optimization**</mark><mark style="color:blue;">:</mark> vLLM optimizes memory efficiency through PagedAttention, which uses concepts from virtual memory and paging to manage KV cache tensors, reducing fragmentation and enabling larger batch sizes.

### <mark style="color:purple;">**Parallelisation Strategies for Large Models**</mark>

* **Necessity for Parallelisation**: Serving models exceeding a single GPU's memory capacity, or meeting strict latency requirements, necessitates parallelization across multiple GPUs.
* **Parallelism Methods**:
  * <mark style="color:blue;">**Tensor Parallelism**</mark><mark style="color:blue;">:</mark> Distributes the computational load of tensor operations across GPUs.
  * <mark style="color:blue;">**Sequence Parallelism**</mark><mark style="color:blue;">:</mark> Focuses on distributing different sequence parts across GPUs.
  * <mark style="color:blue;">**Pipeline Parallelism**</mark><mark style="color:blue;">:</mark> Involves dividing the model into stages, each processed on different GPUs, often in combination with other parallelism methods.

<mark style="color:green;">**Integration of Methods**</mark>

Combining various parallelism techniques (tensor, sequence, pipeline) has been proposed to effectively handle the demands of very large models, ensuring efficient utilization of computational resources and minimizing latency.

In summary, the paper's focus on batching and scheduling highlights S-LoRA's innovative approach to efficiently manage memory and computational resources.&#x20;

This is achieved by separating batched computations for the base model and LoRA adapters, using custom CUDA kernels, and implementing a dynamic iteration-level scheduling strategy.&#x20;

These strategies enable S-LoRA to serve a large number of LoRA adapters concurrently, addressing the challenges of high-throughput and online serving in large language models.

### <mark style="color:purple;">**Batching Strategy**</mark>

* **Goal**: To *<mark style="color:yellow;">support online and high-throughput serving of many LoRA adapters simultaneously</mark>*.
* **Separate Batched Computation**: The base model computation and the LoRA adapters are batched separately. The base computation is batched using General Matrix Multiply (GEMM), while LoRA adapters are handled by custom CUDA kernels.
* **Efficiency**: This separation allows for efficient computation without the need for padding, as would be required by traditional batch GEMM kernels due to heterogeneity in sequence lengths and adapter ranks.

### <mark style="color:purple;">**LoRA Adapter Merging**</mark>

* **Original LoRA Method**: In the original LoRA approach, <mark style="color:yellow;">adapter weights were merged into the base model,</mark> <mark style="color:green;">**creating a new model without additional overhead during inference**</mark>.
* **Limitation in Multiple Adapter Context**: Merging weights for multiple adapters leads to multiple weight copies and misses batching opportunities. Additionally, the original method of adding and subtracting LoRA weights on the fly doesn’t support concurrent inference on separate adapters.

### <mark style="color:purple;">**S-LoRA's Approach**</mark>

* <mark style="color:blue;">**On-the-fly LoRA Computation**</mark><mark style="color:blue;">:</mark> S-LoRA computes the LoRA computation (`xAB`) on-the-fly, avoiding weight duplication and enabling batching of the more costly `xW` operation. This results in considerable savings despite increased computation overhead.
* <mark style="color:blue;">**Custom CUDA Kernels**</mark><mark style="color:blue;">:</mark> Custom CUDA kernels are used for the LoRA computation, optimizing for efficiency and avoiding the limitations of padding required by batch GEMM kernels.

### <mark style="color:purple;">**Memory Management and Adapter Storage**</mark>

* <mark style="color:blue;">**Main Memory Storage**</mark><mark style="color:blue;">:</mark> All LoRA adapters are stored in the main memory. Only the adapters needed for the current batch are fetched to the GPU RAM during inference, which manages the GPU memory more effectively.
* <mark style="color:blue;">**Maximizing Adapter Service**</mark><mark style="color:blue;">:</mark> The maximum number of adapters that can be served is bounded by the main memory size. This approach allows for a higher throughput of adapter serving.

### <mark style="color:purple;">**Iteration-Level Scheduling Batching Strategy**</mark>

* **Adopted from Orca**: The iteration-level scheduling batching strategy from Orca is employed, <mark style="color:yellow;">where requests are scheduled at the token level.</mark>
* **Dynamic Request Incorporation**: New requests are immediately incorporated into the running batch if space is available. Requests exit the batch once they reach the maximum number of generated tokens or other stopping criteria.
* **Memory Management Challenges**: This strategy reduces GPU memory usage but introduces new challenges in memory management.

### <mark style="color:purple;">**Memory Management Techniques**</mark>

* **Efficient Handling of Memory**: Techniques for managing memory efficiently in this dynamic and high-throughput context are discussed in Section 5 of the paper.

In summary, the paper's focus on batching and scheduling highlights S-LoRA's innovative approach to efficiently manage memory and computational resources. This is achieved by <mark style="color:yellow;">separating batched computations for the base model and LoRA adapters, utilizing custom CUDA kernels, and implementing a dynamic iteration-level scheduling strategy.</mark>&#x20;

These strategies enable S-LoRA to serve a large number of LoRA adapters concurrently, addressing the challenges of high-throughput and online serving in large language models.

The section on "Prefetching and Overlapping" and "Custom Kernels" in the S-LoRA paper focuses on optimizing the system's efficiency in handling multiple LoRA adapters. Here's a brief summary:

### <mark style="color:purple;">**Prefetching and Overlapping**</mark>

* <mark style="color:blue;">**Challenge**</mark><mark style="color:blue;">:</mark> I/O overhead from loading and offloading adapter weights, especially with numerous or large adapters, introduces latency.
* <mark style="color:blue;">**Dynamic Prediction Mechanism**</mark><mark style="color:blue;">:</mark> While decoding the current batch, the system predicts which adapters will be needed for the next batch based on the current queue.
* <mark style="color:blue;">**Prefetching Strategy**</mark><mark style="color:blue;">: A</mark>dapters predicted to be needed are prefetched and stored in memory, reducing I/O time and latency for adapter swapping.

### <mark style="color:purple;">**Custom Kernels for Heterogeneous LoRA Batching**</mark>

* <mark style="color:blue;">**Non-Contiguous Memory Handling**</mark><mark style="color:blue;">:</mark> Due to the unified memory pool's design, adapter weights are stored in non-contiguous memory.
* <mark style="color:blue;">**Custom CUDA Kernels**</mark><mark style="color:blue;">:</mark> These kernels are implemented to efficiently handle LoRA computations with varying ranks and sequence lengths in a non-contiguous memory layout.
* <mark style="color:blue;">**Multi-size Batched Gather Matrix-Matrix Multiplication (MBGMM)**</mark><mark style="color:blue;">:</mark> This kernel, implemented in Triton with tiling, is used in the prefill stage to handle a sequence of tokens and gather adapter weights of different ranks.
* <mark style="color:blue;">**Multi-size Batched Gather Matrix-Vector Multiplication (MBGMV)**</mark><mark style="color:blue;">:</mark> Used in the decode stage, this kernel is adapted from Punica to support multiple ranks in a batch and more fine-grained memory gathering for a single token.

### <mark style="color:purple;">**Tensor Parallelism in Batched LoRA Inference**</mark>

* **Objective**: To support multi-GPU inference of large transformer models using batched LoRA inference.
* **Why Tensor Parallelism?**: It's widely used because of its simplicity and compatibility with existing systems. It helps reduce memory usage and latency per GPU when serving large models.

### <mark style="color:purple;">**Partition Strategy**</mark>

* <mark style="color:blue;">**Alignment with Base Model**</mark><mark style="color:blue;">:</mark> The strategy aligns the partitioning of inputs and outputs of the added LoRA computation with those of the base model. This minimizes communication costs by reducing unnecessary communications and fusing some communications.
* <mark style="color:blue;">**Illustration with Feed-Forward Module**</mark>: The partition strategy is demonstrated using a 2-layer MLP (feed-forward module). The base model's first weight matrix (W1) is column-partitioned, and the second (W2) is row-partitioned.
* <mark style="color:blue;">**Added LoRA Computation**</mark><mark style="color:blue;">:</mark> The LoRA computation involves additional matrices (A1, B1, A2, B2). A1 and B1 (for the first weight matrix W1) are column-partitioned, while A2 and B2 (for the second weight W2) are row-partitioned and column-partitioned, respectively.

### <mark style="color:purple;">**Communication in Partitioning**</mark>

* <mark style="color:blue;">**Base Model**</mark><mark style="color:blue;">:</mark> An all-reduce operation accumulates the partial sum from distributed devices.
* <mark style="color:blue;">**LoRA Computation**</mark><mark style="color:blue;">:</mark> Involves all-gather operations to collect intermediate results and all-reduce operations to sum up these results.
* <mark style="color:blue;">**Fusion of Operations**</mark><mark style="color:blue;">:</mark> The strategy includes fusing an all-gather operation with the final all-reduce operation, a novel approach in parallelization.

### <mark style="color:purple;">**Adaptation to Self-Attention Layer**</mark>

* <mark style="color:blue;">**Head Dimension Partitioning**</mark><mark style="color:blue;">:</mark> Similar to the Megatron-LM strategy, the head dimension of the self-attention layer is partitioned.
* <mark style="color:blue;">**Projection Weight Matrices**</mark><mark style="color:blue;">:</mark> The query-key-value projection weight matrix is treated like W1, and the output projection weight matrix is like W2 in the feed-forward module example.

### <mark style="color:purple;">**Communication and Memory Cost Analysis**</mark>

* <mark style="color:blue;">**Costs for Base Model and Added LoRA Computation**</mark><mark style="color:blue;">:</mark> The communication cost for the base model involves one all-reduce operation. For the added LoRA computation, it involves three all-gathers and one all-reduce.
* <mark style="color:blue;">**Negligible Additional Cost**</mark><mark style="color:blue;">:</mark> The extra communication cost introduced by LoRA is small compared to the base model's cost, mainly because the rank of the adapter (r) is much smaller than the hidden size (h).
* <mark style="color:blue;">**Optimal Memory Usage**</mark><mark style="color:blue;">:</mark> The strategy is memory-efficient as it partitions all weight matrices across devices without replication.