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

arXiv: https://arxiv.org/abs/2403.11366 · HTML: https://mlsanity.com/paper/2403.11366/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RAG pipeline design and evaluation, one of 15 directions in Retrieval-augmented generation.
- Builds on 17 of its 26 references that are in the map.
- First posted on arXiv 17 Mar 2024.

## Placed in

- Retrieval-augmented generation / [RAG pipeline design and evaluation](https://mlsanity.com/rag/rag-pipeline-design-and-evaluation/)

## Builds on (most cited first, 17 of 17)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Reading Wikipedia to Answer Open-Domain Questions](https://mlsanity.com/paper/1704.00051/) (2017)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [Compositional Semantic Parsing on Semi-Structured Tables](https://mlsanity.com/paper/1508.00305/) (2015)
- [Internet-Augmented Dialogue Generation](https://mlsanity.com/paper/2107.07566/) (2021)
- [Atlas: End-to-End 3D Scene Reconstruction from Posed Images](https://mlsanity.com/paper/2003.10432/) (2020)
- [Understanding Social Reasoning in Language Models with Language Models](https://mlsanity.com/paper/2306.15448/) (2023)
- [On the Effectiveness of Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2211.15583/) (2022)
- [Simple and Effective Multi-Paragraph Reading Comprehension](https://mlsanity.com/paper/1710.10723/) (2017)
- [Real-Time Open-Domain Question Answering with Dense-Sparse Phrase Index](https://mlsanity.com/paper/1906.05807/) (2019)
- [Raiders of the Lost Kek: 3.5 Years of Augmented 4chan Posts from the Politically Incorrect Board](https://mlsanity.com/paper/2001.07487/) (2020)

## Built on by (most cited first, 2 of 2)

- [Evaluating LLMs Capabilities Towards Understanding Social Dynamics](https://mlsanity.com/paper/2411.13008/) (2024)
- [Advancing Compositional LLM Reasoning with Structured Task Relations in Interactive Multimodal Communications](https://mlsanity.com/paper/2507.21199/) (2025)

## Abstract (arXiv)

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.