JORA：用于检索增强微调的JAX张量并行LoRA库
JORA: JAX Tensor-Parallel LoRA Library for Retrieval Augmented Fine-Tuning
AI总结：
针对RAG任务中LLM微调的内存瓶颈，提出基于JAX张量并行和LoRA的JORA库，实现高效分布式微调，在四GPU上运行速度提升超12倍且显存减半。
摘要：
大型语言模型（LLM）在基于检索的任务（尤其是检索增强生成（RAG））中的扩展面临显著的内存约束，特别是在微调长提示序列时。当前的开源库支持跨多GPU的全模型推理和微调，但无法满足检索上下文所需的高效参数分布。针对这一空白，我们引入了一个新的框架，用于利用分布式训练对Llama-2模型进行PEFT兼容的微调。我们的框架独特地利用JAX的即时（JIT）编译和张量分片来实现高效的资源管理，从而在降低内存需求的同时加速微调。这一进展显著提高了在GPU资源有限的系统上对复杂RAG应用进行LLM微调的可扩展性和可行性。我们的实验表明，与使用四块GPU的Hugging Face/DeepSpeed实现相比，运行时间提升了超过12倍，同时每块GPU的显存消耗不到其一半。
英文摘要：
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 Llama-2 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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
