JORA : JAX Tensor-Parallel L o RA Library for Retrieval Augmented Fine-Tuning (2024)

Abstract: The scaling of Large Language Models (LLMs) for retrieval-based tasks, particularly in Retrieval Augmented Generation (RAG), faces significant memory constraints, especially when fine-tuning extensive prompt sequences. Current open-source libraries support full-model inference and fine-tuning across multiple GPUs but fall short of accommodating the efficient parameter distribution required for retrieved context. Addressing this gap, we introduce a novel framework for PEFT-compatible fine-tuning of GPT models, leveraging distributed training. Our framework uniquely utilizes JAX’s just-in-time (JIT) compilation and tensor-sharding for efficient resource management, thereby enabling accelerated fine-tuning with reduced memory requirements. This advancement significantly improves the scalability and feasibility of fine-tuning LLMs for complex RAG applications, even on systems with limited GPU resources. Our experiments show more than 12x improvement in runtime compared to Hugging Face/DeepSpeed implementation with four GPUs while consuming less than half the VRAM per GPU.

## introduction large language models (llms) like chat-gpt @xcite have revolutionized the field of natural language processing, paving the way for open-source alternatives that offer more flexibility in fine-tuning. llama-2 @xcite , a prominent llm, exemplifies this trend, offering extensive customization at the architecture level. alongside, parameter efficient fine-tuning (peft) @xcite techniques like low-rank adaptation have emerged, optimizing resource utilization in training these models. retrieval augmented generation (rag) @xcite is a paradigm that leverages a corpus to enrich llm prompts with relevant context. however, when fine-tuning on retrieval-based context, the quadratic memory scaling of transformer models with prompt length poses significant challenges, especially when integrating large context sizes. the training process, which employs teacher-forcing at each step of the sequence, exacerbates memory demands, creating a bottleneck for effective llm utilization in rag. current machine learning frameworks facilitate llm fine-tuning on distributed systems, employing model and pipeline parallelism strategies. however, these frameworks lack support for peft, specifically in the context of parallel training. while libraries such as deepspeed @xcite and accelerate @xcite offer data parallelism for fine-tuning the entire model, these libraries lack support for tensor-parallel training in the peft setting. in addition, combining multiple libraries adds unnecessary boilerplate code to glue together dependencies required for parameter-efficient and distributed training. these libraries also require boilerplate code for configuration since they target multiple models. to bridge this gap, we introduce jora (jaxbased lora), a library tailored for llama-2 models, designed to enhance the fine-tuning process for rag applications. utilizing jax's just-in-time (jit) compilation and innovative tensor-sharding techniques, jora not only accelerates the finetuning process but also significantly optimizes memory usage @xcite . our evaluations across standard training gpus demonstrate substantial improvements in training time and memory efficiency, addressing the critical challenges of peft in retrieval-based training. our library also provides valuable helpers for using instruct format datasets, merging lora parameters, and convert-1 152 ing fine-tuned models to hugging face compatible formats. our work makes peft more accessible and efficient for llms, particularly in resourceconstrained environments. by enhancing the scalability and efficiency of llms in retrieval augmented fine-tuning (raft), jora opens new avenues for advanced natural language processing applications. ## jora framework jora is a library for raft. its purpose is to make fine-tuning based on retrieved context more userfriendly. in addition, it is designed to make raft faster and more resource-efficient. figure 1 gives a high-level overview of jora. ## an example usage scenario jora is designed to aid in raft. in this section, we demonstrate a raft use case by fine-tuning it on a social media dataset @xcite to help llms enable social-context understanding. the purpose of rag is to add additional context to a prompt by searching for knowledge and adding additional information. for raft, data can be created based on retrieved knowledge. the llm learns to generate the retrieved answer based on the context since the key rationale is held back. a simple example is a database query, which corresponds to a process that may be taken to produce an output by evaluating the database. if the query is not provided but rather a natural language equivalent is provided, the llm must learn the heuristics represented by the hidden query. since prompt tuning is insufficient for models to develop social-context understanding @xcite , we use a fine-tuning process consisting of two phases to add knowledge to an llm. both phases of fine-tuning use peft. for our problem setting, rather than just predicting the following words, we aim to gain an understanding of the relation across different comments in a social media session. for instance, a comment in a social media session may target the previous comment, the original post that spawned the session, or some comment in the middle of the discourse. to glean insight into the target of the comment in terms of its context, reasoning between the structure of the conversation is critical. unfortunately, the llm pre-training does not consider these relationships specifically, and there is no public data related to reasoning at the comment level in social media discourse. thus, we rely on other general-purpose structured data as a surrogate to learn structure and reasoning. we use the wikitablequestions (pasupat and liang, 2015) dataset to infuse structural intelligence into the model. this dataset consists of various independent tables, questions based on one of the tables, and a corresponding answer. to answer these questions, using the data in the input table is vital. some answers require aggregate reasoning. for the directionality analysis task (which post is targeted by another comment in the same session), we leveraged a corpus of 4chan threads @xcite . this dataset consists of ∼3 million threads and ∼100 million posts. since 4chan allows its users to tag whom they reply to, we use this data as the ground truth for directionality information. we examine whether our raft phases ## conclusion this paper presents jora, a jax-based library for retrieval augment fine-tuning of llama-2 models. jora provides convenient functions for data manipulation and training. in addition, it implements best practices for memory efficient and performant training. by using a combination of lora, 6 157 tensor-parallelism, and jit, jora can significantly improve memory efficiency and computation time over a distributed environment compared to hugging face/deepspeed. finally, jora can export trained models to the popular hugging face model format for downstream usage with other hugging face-compatible libraries.