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

📄 arXiv: 2403.11366v2 📥 PDF

作者: Anique Tahir, Lu Cheng, Huan Liu

分类: cs.LG, cs.CL, cs.DC

发布日期: 2024-03-17 (更新: 2024-03-19)

💡 一句话要点

提出JORA框架以解决大规模语言模型的内存限制问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大规模语言模型 检索增强生成 分布式训练 JAX 张量分片 微调 资源管理

📋 核心要点

  1. 现有方法在处理检索增强生成任务时，面临显著的内存限制，尤其是在微调长提示序列时。
  2. 本文提出了一种基于JAX的分布式训练框架，利用即时编译和张量分片技术，实现高效的参数分配和资源管理。
  3. 实验结果显示，本文方法在运行时间上比现有基线提高了12倍以上，同时显存消耗显著降低。

📝 摘要（中文）

大规模语言模型（LLMs）在检索增强生成（RAG）任务中的扩展面临显著的内存限制，尤其是在微调大量提示序列时。现有的开源库虽然支持多GPU的全模型推理和微调，但在处理检索上下文所需的高效参数分配方面存在不足。为了解决这一问题，本文提出了一种新颖的PEFT兼容的Llama-2模型微调框架，利用分布式训练。该框架独特地采用了JAX的即时编译（JIT）和张量分片技术，以实现高效的资源管理，从而加速微调并减少内存需求。实验结果表明，与Hugging Face/DeepSpeed在四个GPU上的实现相比，运行时间提高了12倍以上，同时每个GPU的显存消耗不到一半。

🔬 方法详解

问题定义：本文旨在解决大规模语言模型在检索增强生成任务中的内存限制问题。现有方法在多GPU环境下无法高效分配参数，导致资源浪费和性能瓶颈。

核心思路：论文提出的框架通过分布式训练和JAX的即时编译技术，优化了模型微调过程中的资源管理，提升了训练效率。

技术框架：整体架构包括数据并行和模型并行两个主要模块，首先将数据分割并分配到多个GPU，然后通过张量分片技术实现模型参数的高效分配。

关键创新：最重要的技术创新在于结合了JAX的即时编译和张量分片，显著提高了微调过程的速度和内存利用率，与传统方法相比，能够在资源有限的情况下实现更高的性能。

关键设计：在参数设置上，采用了动态的张量分片策略，损失函数设计上优化了计算图的构建，确保了在多GPU环境下的高效计算。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，JORA框架在运行时间上相比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.