JORA: JAX Tensor-Parallel LoRA Library for Retrieval Augmented Fine-Tuning

当前，针对基于检索的任务，特别是检索增强生成（RAG）中的大型语言模型（LLMs）的扩展面临着重大的内存限制，特别是在微调大量提示序列时。目前的开源库支持在多个GPU上进行完整模型推理和微调，但无法满足所需的检索上下文的高效参数分布。为了解决这个问题，我们引入了一种新的框架，用于PEFT兼容的Llama-2模型微调，利用分布式训练。我们的框架独特地利用了JAX的即时编译（JIT）和张量分片以实现高效的资源管理，从而实现了加速微调和减少内存需求。这一进展显著提高了复杂RAG应用的LLMs微调的可扩展性和可行性，即使在GPU资源有限的系统上也是如此。我们的实验表明，与使用四个GPU的Hugging Face/DeepSpeed实现相比，运行时间提高了12倍以上，每个GPU消耗的VRAM不到一半。

本论文旨在解决大型语言模型在检索增强生成（RAG）中面临的内存限制问题，特别是在微调大量提示序列时。

本文提出了一种新的框架，利用分布式训练，通过JAX的即时编译和张量分片实现PEFT兼容的Llama-2模型微调，从而实现高效的资源管理，显著提高了大型语言模型微调的可扩展性和可行性。

本文的实验显示，与使用四个GPU的Hugging Face / DeepSpeed实现相比，本文的方法在运行时间上提高了12倍以上，每个GPU消耗的VRAM不到一半。此外，本文还开源了代码。

最近的相关研究包括： 1. Efficient Transformers: A Survey of Learning Paradigms 2. Fine-Tuning Language Models from Human Preferences