SLoRA: Federated Parameter Efficient Fine-Tuning of Language Models
Leveraging Lora
PreviousReLoRA: High-Rank Training Through Low-Rank UpdatesNextGaLora: Memory-Efficient LLM Training by Gradient Low-Rank Projection
Last updated
Copyright Continuum Labs - 2023
Leveraging Lora
Last updated
The paper on S-LoRA focuses on scalable serving of Low-Rank Adaptation (LoRA) adapters for language models (LMs).
Pretrain-then-Finetune Paradigm: LMs are commonly fine-tuned for specific tasks, leading to numerous fine-tuned variants of a single base model.
LoRA for Fine-Tuning: 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.
Challenges in Serving LoRA Adapters: While fine-tuning has been extensively researched, the efficient serving of these fine-tuned variants, especially at scale, remains underexplored.
Goal: To scalably serve thousands of LoRA adapters on a single machine or across multiple GPUs.
Unified Paging: 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.
Tensor Parallelism and CUDA Kernels: It employs a tensor parallelism strategy and custom CUDA kernels for heterogeneous batching of LoRA computation.
Efficiency: S-LoRA aims to serve multiple adapters concurrently with high throughput and low latency, significantly improving over traditional methods.
Adapter and Base Model Separation: S-LoRA separates the batchable base model computation from individual LoRA computations to achieve high-throughput multi-adapter serving.
Memory Management: 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.
Serving Throughput and Latency: Compared to existing libraries (like HuggingFace PEFT and vLLM), S-LoRA can increase throughput by up to 4 times and serve significantly more adapters.
Scalable Serving of Fine-Tuned Models: S-LoRA enables scalable serving of many task-specific fine-tuned models.
Large-Scale Customized Fine-Tuning Services: It offers the potential for large-scale customized fine-tuning services, which could be a substantial advancement for deploying personalized AI solutions.
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.
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.
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.
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.
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.
Autoregressive Decoding: Inference requires iterative autoregressive decoding, where the model encodes the prompt in a forward pass and then decodes the output one token at a time.
Sequential Nature: The sequential nature of decoding each token, attending to the hidden states of all preceding tokens, makes the process slow.
KV Cache: There's a need to store the hidden states of all previous tokens, referred to as the KV cache, which adds substantial memory overhead.
Dynamic Requests: In online settings, handling requests of varying sequence lengths that arrive dynamically is a challenge.
Fine-Grained Scheduling: 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.
vLLM's Memory Optimization: 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.
Necessity for Parallelisation: Serving models exceeding a single GPU's memory capacity, or meeting strict latency requirements, necessitates parallelization across multiple GPUs.
Parallelism Methods:
Tensor Parallelism: Distributes the computational load of tensor operations across GPUs.
Sequence Parallelism: Focuses on distributing different sequence parts across GPUs.
Pipeline Parallelism: Involves dividing the model into stages, each processed on different GPUs, often in combination with other parallelism methods.
Integration of Methods
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.
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.
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.
Goal: To support online and high-throughput serving of many LoRA adapters simultaneously.
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.
Original LoRA Method: In the original LoRA approach, adapter weights were merged into the base model, creating a new model without additional overhead during inference.
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.
On-the-fly LoRA Computation: 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.
Custom CUDA Kernels: Custom CUDA kernels are used for the LoRA computation, optimizing for efficiency and avoiding the limitations of padding required by batch GEMM kernels.
Main Memory Storage: 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.
Maximizing Adapter Service: 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.
Adopted from Orca: The iteration-level scheduling batching strategy from Orca is employed, where requests are scheduled at the token level.
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.
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 separating batched computations for the base model and LoRA adapters, utilizing custom CUDA kernels, and implementing a dynamic iteration-level scheduling strategy.
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:
Challenge: I/O overhead from loading and offloading adapter weights, especially with numerous or large adapters, introduces latency.
Dynamic Prediction Mechanism: While decoding the current batch, the system predicts which adapters will be needed for the next batch based on the current queue.
Prefetching Strategy: Adapters predicted to be needed are prefetched and stored in memory, reducing I/O time and latency for adapter swapping.
Non-Contiguous Memory Handling: Due to the unified memory pool's design, adapter weights are stored in non-contiguous memory.
Custom CUDA Kernels: These kernels are implemented to efficiently handle LoRA computations with varying ranks and sequence lengths in a non-contiguous memory layout.
Multi-size Batched Gather Matrix-Matrix Multiplication (MBGMM): 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.
Multi-size Batched Gather Matrix-Vector Multiplication (MBGMV): 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.
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.
Alignment with Base Model: 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.
Illustration with Feed-Forward Module: 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.
Added LoRA Computation: 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.
Base Model: An all-reduce operation accumulates the partial sum from distributed devices.
LoRA Computation: Involves all-gather operations to collect intermediate results and all-reduce operations to sum up these results.
Fusion of Operations: The strategy includes fusing an all-gather operation with the final all-reduce operation, a novel approach in parallelization.
Head Dimension Partitioning: Similar to the Megatron-LM strategy, the head dimension of the self-attention layer is partitioned.
Projection Weight Matrices: 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.
Costs for Base Model and Added LoRA Computation: 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.
Negligible Additional Cost: 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).
Optimal Memory Usage: The strategy is memory-efficient as it partitions all weight matrices across devices without replication.
