论文状态:已完成

LoRAFusion: Efficient LoRA Fine-Tuning for LLMs

发表:2025/10/01
原文链接PDF 下载
价格:0.100000
价格:0.100000
价格:0.100000
已有 3 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出LoRAFusion,一个高效的低秩适配微调系统,旨在优化大语言模型的微调效率。通过图分裂方法,该系统融合内存密集型操作,减少不必要的内存访问,同时通过自适应批处理算法优化多任务调度,从而实现显著的性能提升,最高可达1.96倍加速,提供了LoRA微调的新方向。

摘要

Low-Rank Adaptation (LoRA) has become the leading Parameter-Efficient Fine-Tuning (PEFT) method for Large Language Models (LLMs), as it significantly reduces GPU memory usage while maintaining competitive fine-tuned model quality on downstream tasks. Despite these benefits, we identify two key inefficiencies in existing LoRA fine-tuning systems. First, they incur substantial runtime overhead due to redundant memory accesses on large activation tensors. Second, they miss the opportunity to concurrently fine-tune multiple independent LoRA adapters that share the same base model on the same set of GPUs. This leads to missed performance gains such as reduced pipeline bubbles, better communication overlap, and improved GPU load balance. To address these issues, we introduce LoRAFusion, an efficient LoRA fine-tuning system for LLMs. At the kernel level, we propose a graph-splitting method that fuses memory-bound operations. This design eliminates unnecessary memory accesses and preserves the performance of compute-bound GEMMs without incurring the cost of recomputation or synchronization. At the scheduling level, LoRAFusion introduces an adaptive batching algorithm for multi-job fine-tuning. It first splits LoRA adapters into groups to intentionally stagger batch execution across jobs, and then solves a bin-packing problem within each group to generate balanced, dependency-aware microbatches. LoRAFusion achieves up to 1.96×1.96\times (1.47×1.47\times on average) end-to-end speedup compared to Megatron-LM, and up to 1.46×1.46\times (1.29×1.29\times on average) improvement over mLoRA, the state-of-the-art multi-LoRA fine-tuning system. Our fused kernel achieves up to 1.39×1.39\times (1.27×1.27\times on average) kernel performance improvement and can directly serve as a plug-and-play replacement in existing LoRA systems. We open-source LoRAFusion at https://github.com/CentML/lorafusion.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

LoRAFusion: Efficient LoRA Fine-Tuning for LLMs

1.2. 作者

Zhanda Zhu, Qidong Su, Yaoyao Ding, Kevin Song, Shang Wang, Gennady Pekhimenko。 所有作者均隶属于多伦多大学、Vector Institute 和 NVIDIA。

1.3. 发表期刊/会议

该论文被接收并计划在 2026年第21届欧洲计算机系统会议 (EUROSYS '26) 上发表,该会议将于2026年4月27日至30日在英国爱丁堡举行。EUROSYS 是计算机系统领域顶级的学术会议之一,享有很高的声誉和影响力。

1.4. 发表年份

2026年 (根据 ACM Reference Format)

1.5. 摘要

低秩适配 (LoRA, Low-Rank Adaptation) 已成为大语言模型 (LLMs, Large Language Models) 最主要的参数高效微调 (PEFT, Parameter-Efficient Fine-Tuning) 方法,因为它显著减少了 GPU 内存使用量,同时在下游任务上保持了具有竞争力的微调模型质量。然而,本文指出现有 LoRA 微调系统存在两个关键效率低下之处:首先,它们由于对大型激活张量的冗余内存访问而产生大量的运行时开销;其次,它们错失了在同一组 GPU 上同时微调共享相同基础模型的多个独立 LoRA 适配器的机会,这导致了性能提升的缺失,例如减少流水线气泡 (pipeline bubbles)、更好的通信重叠和改善的 GPU 负载均衡。

为了解决这些问题,本文提出了 LoRAFusion,一个用于 LLMs 的高效 LoRA 微调系统。在内核层面,LoRAFusion 提出了一种图分裂方法,融合了内存密集型操作。这种设计消除了不必要的内存访问,并在不引入重新计算或同步成本的情况下,保持了计算密集型通用矩阵乘法 (GEMM, General Matrix Multiplication) 的性能。在调度层面,LoRAFusion 引入了一种针对多任务微调的自适应批处理算法。它首先将 LoRA 适配器分成组,以有意地错开批处理执行,然后解决每组内的装箱问题 (bin-packing problem),以生成平衡的、依赖感知的微批次 (microbatches)。

与 Megatron-LM 相比,LoRAFusion 实现了高达 1.96×1.96 \times(平均 1.47×1.47 \times)的端到端加速;与最先进的多 LoRA 微调系统 mLoRA 相比,实现了高达 1.46×1.46 \times(平均 1.29×1.29 \times)的改进。本文提出的融合内核实现了高达 1.39×1.39 \times(平均 1.27×1.27 \times)的内核性能提升,并且可以直接作为现有 LoRA 系统中的即插即用替代品。LoRAFusion 已在 GitHub 上开源。

1.6. 原文链接

2. 整体概括

2.1. 研究背景与动机

核心问题: 尽管 LoRA 已成为大语言模型 (LLMs) 微调的主流参数高效微调 (PEFT) 方法,但现有 LoRA 微调系统仍存在显著的效率瓶颈,未能充分利用 LoRA 的独特特性。

为什么重要: 预训练 LLMs(如 GPT、LLaMa)在文本生成、问答、代码生成等任务中表现出强大能力。为了将这些模型应用于个性化或领域特定任务(如生物医学分析、个性化聊天机器人),微调 (fine-tuning) 是必不可少的。然而,传统的全模型微调需要大量的硬件资源,例如 LLaMa-3.1-70B 的全模型微调仅模型状态就需要约 1120GB 的 GPU 内存,成本高昂且不切实际。LoRA 等 PEFT 方法通过冻结预训练参数并只更新一小部分注入的“适配器”参数,显著降低了资源需求。例如,使用 LoRA 微调 LLaMa-3.1-70B,参数量仅增加 0.29%,GPU 内存使用量降至 142GB。由于这种显著的资源节约,LoRA 微调在云平台和本地环境中被广泛采用。在实际应用中,用户通常会并行运行多个微调任务,以探索不同的超参数或持续适应不断演变的数据集,因此,以每秒处理的训练样本数衡量的微调吞吐量 (throughput) 已成为降低成本和总训练时间的关键指标。

现有研究的挑战或空白: 论文指出现有 LoRA 微调系统存在两个主要的效率问题:

  1. 高运行时开销源于冗余内存访问: 尽管 LoRA 适配器引入的参数量和浮点运算 (FLOPs) 极少,但实际性能分析显示,应用 LoRA 会使训练吞吐量降低约 40%。这主要是因为 LoRA 投影层是内存带宽受限的,并且在适配器操作中重复加载和存储大型激活张量,导致全局内存访问量增加高达 2.64 倍。现有系统,如 PEFT Library、LLaMA-Factory 等,主要依赖于为全模型微调设计的并行化方法(如 FSDP、TP、PP),这些方法未能充分解决 LoRA 自身的特性所带来的效率问题。
  2. 错失多任务(Multi-LoRA)并发微调机会: 现有系统未能有效支持同时微调多个共享相同基础模型的独立 LoRA 适配器。每个适配器通常独立微调,即使它们足够轻量化可以合并。在超参数调优或多租户云服务等场景中,并发微调多个 LoRA 适配器(Multi-LoRA fine-tuning)具有巨大潜力。虽然服务场景中广泛使用 Multi-LoRA 优化来提高 GPU 利用率,但在微调场景中,其主要优势在于缓解分布式并行带来的开销(如减少流水线气泡、改善通信重叠、平衡 GPU 负载)。现有 Multi-LoRA 微调系统 (如 mLoRA) 仍有局限性,例如依赖通用 LoRA 内核(受冗余内存访问限制),未能有效解决 GPU 间的负载不平衡问题(尤其是在可变序列长度场景下)。

论文的切入点/创新思路: LoRAFusion 认为,一个高效的 LoRA 微调系统必须同时减少 LoRA 适配器的运行时开销,并利用多 LoRA 优化机会来减少分布式训练开销。为此,它提出了一个多层级融合系统,在内核层面和调度层面进行优化:

  • 内核层面: 通过图分裂方法融合内存密集型操作,消除冗余内存访问。
  • 调度层面: 引入自适应批处理算法,通过适配器分组和装箱问题解决,实现多任务并发微调的负载均衡和流水线优化。

2.2. 核心贡献/主要发现

论文最主要的贡献:

  1. 识别并解决核心瓶颈: 首次明确指出并解决了现有 LoRA 微调系统中的两大关键限制:冗余内存访问导致的高运行时开销,以及未能充分利用多任务并发微调带来的优化机会。
  2. 提出多层级融合系统 LoRAFusion:
    • 内核层面的水平融合策略: 设计了 FusedLoRAFusedMultiLoRA 内核,通过在计算图中战略性地分裂和融合操作,减少内存带宽瓶颈,同时保持计算密集型操作的性能。
    • 调度层面的多 LoRA 调度算法: 引入了创新的自适应批处理算法,通过适配器分组和基于混合整数线性规划 (MILP, Mixed Integer Linear Programming) 的装箱问题解决,优化了多任务并发微调的 GPU 负载均衡和分布式训练开销(如流水线气泡)。
  3. 广泛的实验验证: 在 LLaMa3.1-8B、Qwen-2.5-32B 和 LLaMa-3.1-70B 等多种 LLMs、真实数据集和 GPU 平台上(NVIDIA H100 和 L40S)进行了广泛评估。
  4. 提供可插拔的融合内核: FusedLoRA 内核可以作为现有 LoRA 系统的即插即用替代品,为社区带来即时性能提升。

论文得出的关键结论或发现:

  • LoRAFusion 相比 Megatron-LM 实现了高达 1.96×1.96 \times(平均 1.47×1.47 \times)的端到端加速。

  • LoRAFusion 相比最先进的 Multi-LoRA 微调系统 mLoRA 实现了高达 1.46×1.46 \times(平均 1.29×1.29 \times)的改进。

  • 融合内核 FusedLoRA 自身就实现了高达 1.39×1.39 \times(平均 1.27×1.27 \times)的性能提升。

  • 通过融合,GPU DRAM 内存读/写流量减少了 34%37%34\%-37\%

  • 通过调度,流水线气泡率从 4 个适配器时的 mLoRA 的 34.11%34.11\% 显著降低到 LoRAFusion 的 11.09%11.09\%,将 GPU 利用率从 65%65\% 提高到 89%89\%

  • LoRAFusion 在不同的模型规模、硬件平台和数据集配置下均表现出强大的通用性和鲁棒性,尤其是在异构(Heterogeneous)数据设置下。

  • 在多 GPU 场景下,任务级(job-level)扩展(运行更多并发任务)的效率高于数据并行(DP)扩展(单个任务使用更多 GPU)。


3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 大语言模型 (LLMs, Large Language Models)

大语言模型是拥有数亿到数万亿参数的深度神经网络模型,通过在海量文本数据上进行预训练而获得强大的语言理解和生成能力。它们能够执行文本生成、问答、翻译、摘要等多种自然语言处理任务。例如,GPT 和 LLaMa 系列模型。

3.1.2. 微调 (Fine-tuning)

微调是指在特定任务或领域的数据集上,对一个已经预训练好的模型进行进一步训练的过程。其目的是使模型能够更好地适应特定的下游任务,同时保留其在预训练阶段学到的通用知识。传统的全模型微调需要更新模型的所有参数。

3.1.3. 参数高效微调 (PEFT, Parameter-Efficient Fine-Tuning)

PEFT 是一类优化方法,旨在减少微调大型模型时所需的计算和内存资源。这些方法通过冻结预训练模型的大部分参数,只训练一小部分新引入的或选定的参数(通常称为“适配器”),从而大大降低了微调成本。

3.1.4. 低秩适配 (LoRA, Low-Rank Adaptation)

LoRA 是当前最流行的 PEFT 方法之一。其核心思想是,在预训练模型的每个权重矩阵旁边,注入两个小的、可训练的低秩矩阵来近似权重更新。 对于一个预训练权重矩阵 WRk×nW \in \mathbb{R}^{k \times n},LoRA 引入两个低秩矩阵 ARk×rA \in \mathbb{R}^{k \times r}BRr×nB \in \mathbb{R}^{r \times n},其中秩 rmin(n,k)r \ll \min(n, k)。原始输出 XW 被修改为 XW+α(XA)BXW + \alpha(XA)B。 其数学表示为: Y=XW+αSB=XW+α(X^A)B Y = X W + \alpha S B = X W + \alpha ( \widehat { X } A ) B 其中:

  • YRm×nY \in \mathbb{R}^{m \times n} 是输出张量。

  • XRm×kX \in \mathbb{R}^{m \times k} 是输入张量。

  • WRk×nW \in \mathbb{R}^{k \times n} 是基础模型的预训练权重矩阵,在 LoRA 微调中被冻结(不可训练)。

  • X^Rm×k\widehat{X} \in \mathbb{R}^{m \times k} 是经过 Dropout 操作后的输入张量 XX

  • ARk×rA \in \mathbb{R}^{k \times r} 是第一个 LoRA 权重矩阵(下投影层),是可训练的。

  • BRr×nB \in \mathbb{R}^{r \times n} 是第二个 LoRA 权重矩阵(上投影层),是可训练的。

  • S=X^AS = \widehat{X}A 是中间结果张量。

  • mm 是聚合的词元数(批处理大小乘以序列长度)。

  • kk 是输入维度。

  • nn 是输出维度。

  • rr 是 LoRA 的秩,通常远小于 kknn

  • α\alpha 是一个常数标量,用于缩放 LoRA 分支的输出。

    LoRA 显著减少了微调所需的内存和计算资源,因为它只训练 AABB 矩阵,而这些矩阵的参数数量远少于 WW

LoRA 微调中的符号约定 (Table 1): 以下是本文中 LoRA 微调的符号约定:

符号 描述
r LoRA 秩
m 词元数(批处理大小 X 序列长度)
k,n 权重矩阵的输入/输出维度
W 基础模型权重。大小: (k, n)。
A 第一个 LoRA 权重。大小: (k, r)
B 第二个 LoRA 权重。大小: (r, n)
X 输入张量。大小: (m, k)
X^\widehat{X} Dropout 后的输入张量。大小: (m, k)
Y 输出张量。大小: (m, n)

3.1.5. 并行化 (Parallelization)

为了训练大型模型,通常需要将计算和内存分布到多个 GPU 上:

  • 数据并行 (DP, Data Parallelism): 在每个 GPU 上复制模型,并将输入数据批次分割到各个 GPU 上并行处理。
  • 完全分片数据并行 (FSDP, Fully Sharded Data Parallelism) 或 ZeRO-3: 将模型状态(参数、梯度、优化器状态)分片到不同的 GPU 上,只在需要时进行通信,显著减少内存使用。
  • 张量并行 (TP, Tensor Parallelism): 将线性层(如矩阵乘法)分割到多个 GPU 上,每个 GPU 计算一部分,然后通过通信合并结果。
  • 流水线并行 (PP, Pipeline Parallelism): 将模型的不同层或阶段分配给不同的 GPU,形成一个流水线。这可以减少通信开销,但可能引入空闲时间(称为流水线气泡,pipeline bubbles)。

3.1.6. 内核融合 (Kernel Fusion)

内核融合是一种优化技术,通过将多个 GPU 操作合并到更少的 GPU 内核中,以减少内存传输和内核启动开销,从而提高效率。例如,将一系列连续的逐元素操作融合到一个内核中,可以避免中间结果写入全局内存再读回的开销。

3.1.7. 即时数据打包 (On-the-fly Data Packing)

在 LLM 微调中,训练数据通常包含长度可变的序列。

  • 传统填充 (Traditional Padding, Figure 2a): 将较短的序列用填充词元补齐到批次中最长序列的长度,导致计算浪费。

  • 数据集预打包 (Dataset Pre-packing, Figure 2b): 预先形成固定长度的批次,但可能导致每个批次中的样本数量可变,若处理不当可能影响训练稳定性和随机性。

  • 即时打包 (On-the-fly Packing, Figure 2c): 在每个批次内动态地将多个样本拼接起来,以避免计算浪费,同时保持每个批次中训练样本数量的确定性。本文采用了这种方法。

    Figure 2. Comparison of (a) traditional batch padding, (b) dataset pre-packing, and (c) batch on-the-fly packing methods for LoRA fine-tuning of LLMs. 该图像是图表,展示了三种不同的LoRA微调批处理方法的比较:传统批处理填充(a)、数据集预打包(b)和即时批打包(c)。每种方法的特点及其在处理不平衡批次问题上的表现被具体展示。

以下是原文 Figure 2 的结果: Figure 2 展示了 LoRA 微调 LLM 的三种批处理方法的比较:(a) 传统批处理填充,(b) 数据集预打包,和 (c) 即时批打包。传统填充通过添加填充词元使所有序列达到最大长度,导致计算浪费。数据集预打包在预处理阶段将短序列拼接成固定长度的批次,但可能影响训练随机性。即时打包在运行时动态拼接序列,以优化计算效率同时保持批次内样本数量的稳定性,是本文采用的方法。

3.2. 前人工作与技术演进

3.2.1. LoRA 算法及其变体

LoRA [32] 作为基础,其变体如 DoRA [52] 和 VeRA [37] 也在不断发展,它们通常在核心 LoRA 计算前后增加预处理或后处理功能,以进一步提升性能或效率。这些算法层面的进步不断优化 LoRA 的效果。

3.2.2. LoRA 服务系统

在 LLM 服务场景中,为了处理高并发请求和提高 GPU 利用率,已经出现了多种 Multi-LoRA 技术。例如,PetS [110] 提出了多任务参数高效微调 Transformer 服务,并引入调度算法协调不同请求。Punica [9]、S-LoRA [79] 和 dLoRA [94] 都旨在通过同时服务多个 LoRA 适配器来提高系统吞吐量,以缓解小词元数在解码过程中导致的内存带宽受限问题。

3.2.3. LoRA 微调系统

相比服务场景,LoRA 微调的研究更多集中在隐私保护方面。例如,Offsite-tuning [96] 和 DLoRA [26] 旨在通过轻量级适配器连接大型模型所有者和数据所有者,以增强隐私。在并发工作中,LobRA [47] 探索了在异构数据上处理多租户微调的问题。然而,现有 LoRA 微调系统(如 PEFT Library [55]、LLaMA-Factory [108]、llama-cookbook [57])主要还是复用为全模型微调设计的并行化和优化技术。mLoRA [98] 是一个值得注意的例外,它初步使用了 Multi-LoRA 分组来减少预训练模型的内存占用和提高训练效率,特别是在流水线并行中,但其设计仍有局限性。

3.2.4. 通用系统优化

  • 并行化: 从早期的数据并行 [1, 42] 到 FSDP [73, 104]、张量并行 [53, 81] 和流水线并行 [21, 33, 62],并行化技术不断发展以应对大模型训练的内存和计算需求。混合并行 [7, 81, 106] 结合多种策略以达到最佳性能。
  • 内核融合: Flash-Attention [13] 和元素级内核融合 [31, 84] 等技术通过合并操作来减少内存传输和内核启动开销,显著提高了 LLM 的性能。
  • 数据打包: 为了应对可变长度序列带来的计算浪费,即时数据打包 [38, 90] 动态拼接序列,避免了填充,提高了效率。

3.3. 差异化分析

本文的工作与上述前人工作的主要区别和创新点在于:

  1. 专注 LoRA 微调的独特瓶颈: 现有工作大多针对全模型微调或 LoRA 服务场景进行优化,但 LoRAFusion 明确指出并解决了 LoRA 微调自身特有的两个主要低效问题:

    • LoRA 模块的运行时开销: 现有通用 LoRA 内核(包括 mLoRA 使用的内核)未能解决由于低秩特性导致的内存带宽瓶颈和对大型激活张量的冗余内存访问。LoRAFusion 通过其 FusedLoRAFusedMultiLoRA 内核直接针对此问题进行优化。
    • 多任务微调的调度优化: 现有系统未能有效利用多任务并发微调的机会。mLoRA [98] 虽尝试进行 Multi-LoRA 微调,但存在以下局限性:
      • 不处理负载不平衡: mLoRA 假设适配器分组均匀,未解决真实工作负载中可变序列长度导致的 GPU 负载不平衡问题。而 LoRAFusion 的调度器专门通过自适应批处理来解决此问题。
      • 通用内核性能: mLoRA 的 BatchLoRA 内核主要减少内核启动开销,但未能解决核心的内存冗余访问瓶颈,因此其性能受限于通用 LoRA 内核。LoRAFusion 则提供了优化的融合内核。
      • 并行模式和通信限制: mLoRA 的设计主要针对流水线并行,并依赖效率较低的 CPU 到 CPU 的 RPC 通信,限制了在现代 GPU 集群(如 NVLink 互联)上的可扩展性。LoRAFusion 则基于 Megatron-LM,支持高效的分布式并行和高带宽互联。
  2. 多层级融合策略: LoRAFusion 结合了内核层面的内存优化(减少内存访问)和调度层面的任务优化(减少分布式开销、平衡负载),形成了一个更全面的解决方案。这使得 LoRAFusion 能够协同解决微调效率的多个维度问题。

  3. 流水线感知调度: LoRAFusion 的调度算法特别考虑了流水线并行中的数据依赖性(bubble lemma),通过智能分组和装箱问题,在保证正确性的前提下,最大程度地减少了流水线气泡,这是 mLoRA 等系统所未能充分解决的。


4. 方法论

LoRAFusion 旨在解决 LoRA 微调系统中的两大效率瓶颈:LoRA 模块冗余内存访问导致的运行时开销,以及多任务微调中负载不平衡和并行开销导致的吞吐量下降。为此,它提出了一个双管齐下的解决方案:在内核层面通过融合操作优化内存访问,在调度层面通过智能批处理优化多任务并发。

4.1. FusedLoRA 和 FusedMultiLoRA

4.1.1. 问题背景与设计考量

正如引言所述,LoRA 模块虽然参数量小,但实际运行时开销大。这主要是因为 LoRA 操作中包含大量对大型激活张量的冗余内存访问。为了解决这个问题,内核融合 (kernel fusion) 是一种有效手段,即将多个 GPU 操作合并为一个内核,以减少内存传输和内核启动开销。

然而,将所有 LoRA 操作简单地融合到一个内核中会面临挑战:

  1. 计算密集型 GEMM 的性能敏感性: 基础模型的 GEMM 操作 Y1=XWY_1 = XW 是计算密集型 (compute-bound) 的。它的性能高度依赖于核函数的分块策略 (tiling strategies) 和 GPU 资源(如共享内存和寄存器文件)的有效利用。不佳的融合可能导致次优的分块布局或资源过度使用,从而严重降低 XW 的性能。

  2. 重新计算与同步开销: 融合具有生产者-消费者依赖关系的操作(例如,先计算 X^A\widehat{X}A,再计算 (X^A)B(\widehat{X}A)B)时,可能需要重新计算中间结果或引入线程块 (thread block) 间的同步。这两者都会增加开销。

    为了克服这些挑战,LoRAFusion 引入了一种图分裂策略 (graph-splitting strategy)。其核心思想是在计算图中将大型激活张量缩小到低秩维度 rr 的中间张量处进行分裂,而不是进行完全融合。这样做的好处是:

  • 中间张量小且廉价: 中间张量 S=X^AS = \widehat{X}A 的大小与 LoRA 秩 rr 相关,通常非常小。存储和重新加载它从 GPU 全局内存的成本很低。
  • 避免重新计算和同步: 通过存储中间结果 SS,可以避免在不同部分反复重新计算它,也避免了线程块间复杂的同步机制。
  • 保护计算密集型 GEMM 性能: 这种分裂策略使得计算密集型 XW 操作能够保持其独立性,从而可以采用最优的分块策略和资源利用,确保其峰值性能不受影响。

4.1.2. FusedLoRA 设计

FusedLoRA 内核的设计目标是在 LoRA 计算图中减少内存传输,同时保持计算效率。其前向和反向传播的融合策略如图 9 和图 10 所示。

Figure 9. Overview of our fusion strategy for LoRA modules in the forward pass, illustrating the full graph fusion approach vs. the split graph fusion approach. 该图像是展示 LoRA 模块在前向传递中的融合策略的示意图,比较了完全图融合方法与分裂图融合方法。图中展示了不同选项,包括融合与重计算、融合与同步等操作,以优化内存访问和计算效率。

以下是原文 Figure 9 的结果: Figure 9 概述了 LoRA 模块前向传播的融合策略,比较了全图融合方法与分裂图融合方法。全图融合可能需要在每个 tile 中重复计算 S,或者通过线程块同步共享 S,这两种方式都有开销。LoRAFusion 采用分裂图融合,通过显式存储和重新加载 S 来避免重计算和同步,因为 S 相对较小,且其读写开销较低。

该图像是示意图,展示了LoRAFusion中的前向图和反向图(图 1(a) 和 1(b))以及优化机会与融合内核策略(图 1(c))。图中包含运算符及其关系,突显了内存重载和操作融合的优化。 该图像是示意图,展示了LoRAFusion中的前向图和反向图(图 1(a) 和 1(b))以及优化机会与融合内核策略(图 1(c))。图中包含运算符及其关系,突显了内存重载和操作融合的优化。

以下是原文 Figure 10 的结果: Figure 10 描绘了 LoRAFusion 的内核设计。图 10(a) 展示了前向传播的融合策略:操作 D\mathbb{D} 融合了 Dropout 和下投影 (X^A)(\widehat{X}A),操作 E\mathbb{E} 将基础模型 GEMM (XW) 与 LoRA 上投影 (αSB)(\alpha SB) 融合,直接累加结果。图 10(b) 展示了反向传播的融合策略:操作 I\mathbb{I} 融合了 SSBB 的梯度计算,操作 K\mathbb{K} 保持独立(处理小型掩码输入),操作 L\mathbb{L} 水平融合了基础模型的梯度计算与 LoRA 路径的操作。这种设计通过减少内存访问来提高效率。

1. 前向传播 (Forward Pass, Figure 10a):

  • 操作 D\mathbb{D} (Dropout & Down-Projection Fusion):
    • Dropout 操作和 LoRA 的下投影操作 (X^A)(\widehat{X}A) 融合到一个内核中。
    • 目的: 消除重新加载完整大小激活张量 X^\widehat{X} 的需要。在原始实现中,Dropout 之后可能需要将 X^\widehat{X} 写入内存,然后 X^\widehat{X} 再从内存中读取进行 AA 矩阵乘法。融合后,X^\widehat{X} 可以直接在寄存器或共享内存中进行 Dropout,并立即用于与 AA 矩阵的乘法,减少了内存往返。
  • 操作 E\mathbb{E} (Base GEMM & LoRA Up-Projection Fusion):
    • 将计算密集型基础模型 GEMM (Y1=XWY_1 = XW) 与内存密集型 LoRA 上投影操作 (Y2=αSBY_2 = \alpha SB) 融合。
    • 目的: 消除冗余内存操作,并通过直接累加部分结果来节省对完整大小输出张量的一次读写。在传统实现中,XW 的结果 Y1Y_1 会写入内存,αSB\alpha SB 的结果 Y2Y_2 也会写入内存,然后两者再从内存中读取并相加。融合后,可以在计算 Y1Y_1 的同时或之后立即计算 Y2Y_2,并将两者直接在寄存器或共享内存中相加,最终只将最终结果 Y=Y1+Y2Y = Y_1 + Y_2 写入内存。这在不影响基础 GEMM 性能的前提下,有效减少了内存带宽消耗。

2. 反向传播 (Backward Pass, Figure 10b):

  • 操作 I\mathbb{I} (Gradient of S & B Fusion):
    • 融合了中间张量 SS 和 LoRA 权重 BB 的梯度计算。
    • 目的: 避免重新加载梯度 dY。在计算 dB (即 STdYS^T dY) 和 dS (即 dYBTdY B^T) 时,如果 dY 只需要加载一次,就可以避免重复内存访问。
  • 操作 K\mathbb{K} (Masked Input Gradient):
    • 此操作保持独立,因为它作用于小型掩码输入(masked input)。
    • 目的: 对于这种小规模操作,融合带来的收益很小,且可能引入不必要的复杂性。
  • 操作 L\mathbb{L} (Base Model Gradient & LoRA Path Gradient Fusion):
    • 水平融合了计算密集型基础模型梯度计算和内存密集型 LoRA 路径操作。

    • 目的: 防止对部分输出梯度进行冗余读写。类似于前向传播的 E\mathbb{E} 操作,通过直接累加来自不同路径的梯度,减少对完整大小梯度张量的内存操作。

      核心洞察: FusedLoRA 的关键在于战略性地识别那些共享大型张量(如 D,E\mathbb{D}, \mathbb{E}L\mathbb{L} 涉及的操作)的操作,进行水平融合以减少内存流量,同时不影响计算密集型操作的最优分块策略,从而有效地缓解内存带宽瓶颈。

4.1.3. 扩展至 FusedMultiLoRA

为了支持并发微调多个 LoRA 适配器,FusedLoRA 被扩展为 FusedMultiLoRA。这使得融合内核能够处理来自不同任务的混合适配器批次。

  • 概念: 在同一个批次中,不同的词元 (tokens) 可能属于不同的微调任务,因此需要应用不同的 LoRA 适配器。

  • 机制 (Figure 11):

    • 标记和查找表: 每个输入 Mtile (一个词元块) 都被标记一个适配器 ID 和其配置(如 LoRA 秩、缩放因子 α\alpha、Dropout 比率等)。这些配置存储在一个轻量级的查找表 (lookup table) 中。
    • 运行时路由: 在内核执行期间,基础模型(冻结部分)的计算在所有词元间共享。但是,对于 LoRA 适配器特定的逻辑,会根据 Mtile 的标签动态应用。对于输出的每个 (Mtile, Ntile) 块,内核会从查找表中获取正确的适配器 ID,然后加载相应的 AABB 矩阵,并应用正确的缩放和 Dropout。
    • 反向传播: 在反向传播过程中,使用相同的映射机制将梯度路由到各自的适配器,而不会相互干扰。
    • 效率: 这种瓦片级 (tile-level) 路由允许在单个融合内核运行中高效执行异构适配器,避免了每个适配器单独启动内核的冗余计算,并实现了作业级优化。
  • 动态选择: 系统会根据批次中是否存在多个适配器,动态选择使用 FusedLoRA (单个适配器) 还是 FusedMultiLoRA (多个适配器)。

    Figure 11. Illustration of FusedMultiLoRA in the forward pass. The routing of LoRA adapters is done at the tile level. 该图像是多种 LoRA 适配器的前向传播示意图,展示了 MultiLoRA A 矩阵和 B 矩阵的路由关系。映射查找表在其中起到关键作用,确保适配器间的信息传递与高效调度。

以下是原文 Figure 11 的结果: Figure 11 展示了 FusedMultiLoRA 在前向传播中的运作方式,特别是 LoRA 适配器的瓦片级路由。每个输入的 Mtile 都带有一个适配器 ID,这个 ID 指向一个查找表,该表包含 LoRA 秩、缩放因子和 Dropout 比率等配置信息。在计算过程中,基础模型的计算是共享的,但 LoRA 适配器的 AABB 矩阵、缩放和 Dropout 会根据每个 Mtile 的适配器 ID 动态选择和应用。

4.2. Multi-LoRA 调度器

LoRAFusion 不仅通过融合内核减少运行时开销,还通过调度多个 LoRA 微调任务来提高端到端吞吐量。这通过智能的适配器分组和样本批处理实现,以平衡 GPU 负载并最小化分布式并行开销。

4.2.1. 调度器工作流 (Figure 12)

Figure 12. Multi-LoRA adapter scheduling workflow. Top: Adapter grouping by sequence length statistics. Middle: Twostage MILP optimization for microbatch creation. Bottom: Cross-batch merging of underfilled microbatches. 该图像是示意图,展示了多LoRA适配器的调度工作流程。顶部展示了根据序列长度统计的适配器分组,中间部分显示了微批次创建的双阶段MILP优化,而底部则展示了对不足填充微批次的跨批次合并。

以下是原文 Figure 12 的结果: Figure 12 描绘了 Multi-LoRA 适配器调度器的工作流程。顶部展示了根据序列长度统计进行的适配器分组。中间部分显示了微批次创建的双阶段 MILP 优化过程。底部则说明了在满足数据依赖性的前提下,如何对不足填充的微批次进行跨批次合并。

调度器的工作流程如下:

  1. 输入: 一组微调任务。
  2. 数据统计: 从数据集中提取样本长度统计信息。
  3. 并行性模拟: 通过并行性模拟器确定微批次的词元预算。
  4. 适配器分组 (Adapter Grouping): 根据样本长度分布对 LoRA 适配器进行分组。
  5. 微批次构建 (Microbatch Construction): 对每组内的样本,使用两阶段 MILP 优化算法构建平衡的微批次。
  6. 迭代优化: 对分组和批处理结果进行模拟评估,并迭代调整直至找到高吞吐量配置。
  7. 任务执行: 使用 FusedMultiLoRA 内核执行任务。
  8. 运行时协调: 多适配器运行时协调器确保词元与适配器的一致性,管理资源共享,并跟踪跨任务边界的梯度。

4.2.2. 粒度 (Granularity)

调度器以单个全局批次 (global batch) 的粒度进行操作。每个适配器的数据集被划分为全局批次。然后,调度器将所有属于相同全局批次索引的样本(来自不同适配器)聚合起来,并将其打包成多个微批次。

4.2.3. 气泡引理与适配器分组 (Bubble Lemma & Adapter Grouping)

  • 问题: 在流水线并行中,由于数据依赖性,一个样本的反向传播只有在 S-1 个其他微批次完成前向传播后才能开始。如果来自同一适配器的连续批次样本被分散到不同的微批次中,可能会导致依赖冲突,即批次 jj 的尾部与批次 j+1j+1 的头部发生冲突,从而引发错误的执行。
  • 气泡引理 (Bubble Lemma): 对于适配器 ii,如果来自全局批次 jj 的样本 ss 被提交到微批次 kk,那么来自同一适配器的批次 j+1j+1 的任何样本都不能在微批次 k+S-1 之前被提交(即在样本 ss 的反向传播完成后)。
  • 解决方案: LoRAFusion 通过适配器分组来解决这个冲突。它将适配器分成组,在组之间保持严格的顺序,而在组内允许灵活的合并。这确保了来自同一适配器的批次在调度中被其他组的微批次充分间隔开,从而满足气泡条件。
  • 负载均衡: 为了在组内实现负载均衡,调度器采用头尾配对 (head-tail pairing) 策略:根据平均词元长度对适配器进行排序,并将短序列适配器与长序列适配器配对。

4.2.4. 两阶段 MILP 数据批处理 (Two-Stage MILP Data Batching)

在适配器分组后,调度器需要解决一个装箱问题 (bin-packing problem),将样本打包成微批次,每个微批次都受固定词元容量的约束。目标是:

  1. 最小化所需的微批次总数。

  2. 使最小的微批次尽可能空,以便在后续阶段进行更好的合并。

    该装箱问题通过两阶段的混合整数线性规划 (MILP) 来解决 (见 Algorithm 1,第 3-7 行)。

Algorithm 1: 数据批处理与合并(每组)

数据: 带有分组样本的适配器,词元容量 C,超时 t
结果: 满足流水线约束的调度微批次

1 foreach global batch b in parallel do
2     (Bg, {m}) ← GreedyPacking(b, C) // 作为基线的贪婪回退
3     B* ← MILP_MinBins(b, C, timeout = t) // 阶段 1: 最小化微批次数量
4     if B* ≥ Bg then
5         B* ← Bg
6     end
7     {B1 ← MILP_MinSmallestBin(b, B*, C, timeout = t)} // 阶段 2: 最小化最小箱中的词元数
8     if B* = Bg and m* ≥ m_greedy then // 如果 MILP 没有改进,则使用贪婪结果
9         return GreedyPacking(b,C)
10    end
11 end
12 foreach consecutive batch pairs (b, b+1) do
13     Shift tokens from b+1 into b if bubble lemma is preserved
14 end
15 VerifyAndFix(schedule) // 在需要的地方插入空操作 (no-ops)
16 return Scheduled microbatches

为便于表述,我们定义 PP 为填充倍数(用户指定,例如 64 或 128),用于将属于同一适配器的序列长度填充到 PP 的倍数。

  • xs,b{0,1}x_{s,b} \in \{0, 1\} 表示样本 ss 是否被分配到箱 bb
  • ka,bNk_{a,b} \in \mathbb{N} 表示适配器 aa 在箱 bb 中贡献的填充倍数。
  • zb{0,1}z_b \in \{0, 1\} 表示箱 bb 是否被使用。

1. 第一阶段 (最小化微批次数量): 优化目标是最小化使用的箱子总数 b=1Bzb\sum_{b=1}^{B} z_bargminxs,hs,bs,zb,b=1Bzbs.t.zb+1zbb<Bb=1Bxs,b=1ssamplesstanglersadyrc(a)len(s)xs,bka,bPa,bzbaka,bPcapacityzbb \begin{array} { c l } { \displaystyle \operatorname * { a r g m i n } _ { x _ { s } , h _ { s } , b _ { s } , z _ { b } , } } & { \displaystyle \sum _ { b = 1 } ^ { B } z _ { b } } \\ { \mathrm { s . t . } } & { \displaystyle z _ { b + 1 } \leq z _ { b } \quad \forall b < B } \\ & { \displaystyle \sum _ { b = 1 } ^ { B } x _ { s , b } = 1 \quad \forall s \in \mathrm { s a m p l e s } } \\ & { \displaystyle \operatorname* { s t a n g l e r } _ { s \in \mathrm { a d y r c } ( a ) } \mathrm { l e n } ( s ) \cdot x _ { s , b } \leq k _ { a , b } \cdot P \quad \forall a , b } \\ & { \displaystyle z _ { b } \leq \sum _ { a } k _ { a , b } \cdot P \leq \mathrm { c a p a c i t y } \cdot z _ { b } \quad \forall b } \end{array} 其中:

  • 目标函数: b=1Bzb\sum_{b=1}^{B} z_b 表示使用的箱子总数。
  • 约束条件:
    • zb+1zbb<Bz_{b+1} \leq z_b \quad \forall b < B: 确保使用的箱子是从第一个开始连续的。
    • b=1Bxs,b=1ssamples\sum_{b=1}^{B} x_{s,b} = 1 \quad \forall s \in \mathrm{samples}: 每个样本必须且只能被分配到一个箱子。
    • sadapter(a)len(s)xs,bka,bPa,b\sum_{s \in \mathrm{adapter}(a)} \mathrm{len}(s) \cdot x_{s,b} \leq k_{a,b} \cdot P \quad \forall a, b: 适配器 aa 在箱子 bb 中贡献的总词元数(经过填充)不能超过 ka,bPk_{a,b} \cdot P
    • zbaka,bPcapacityzbbz_b \leq \sum_a k_{a,b} \cdot P \leq \mathrm{capacity} \cdot z_b \quad \forall b: 如果箱子 bb 被使用 (zb=1z_b=1),则其中总词元数不能超过容量 capacity;如果未被使用 (zb=0z_b=0),则其中词元数为 0。

2. 第二阶段 (最小化最小箱中的词元数): 在确定了最优的箱子数量 BB^* 后,第二阶段将 B=BB = B^* 固定,目标是最小化所有箱子中词元总数最小的那个箱子的大小。 arcsinminxs,b,ka,bminb[1,Bs]aka,bPs.t.b=1Bxs,b=1ssamplessadapter(a)len(s)xs,bka,bPa,baka,bPcapacityb \begin{array} { l } { { \displaystyle \arcsin \operatorname* { m i n } _ { \boldsymbol { x } _ { s , b } , \boldsymbol { k } _ { a , b } } \quad \underset { b \in [ 1 , B ^ { s } ] } { \operatorname* { m i n } } \sum _ { a } k _ { a , b } \cdot P } } \\ { { \mathrm { s . t . } \quad \quad \displaystyle \sum _ { b = 1 } ^ { B ^ { * } } x _ { s ,b } = 1 \quad \forall s \in \mathrm { s a m p l e s } } } \\ { { \displaystyle \sum _ { s \in \mathrm { a d a p t e r } ( a ) } \mathrm { l e n } ( s ) \cdot x _ { s ,b } \leq k _ { a ,b } \cdot P \quad \forall a , b } } \\ { { \displaystyle \sum _ { a } k _ { a ,b } \cdot P \leq \mathrm { c a p a c i t y } \quad \forall b } } \end{array} 其中:

  • 目标函数: 最小化所有箱子中词元总数最小的那个箱子的词元数。这使得最不饱满的微批次留有更多空间,以便在后续阶段进行合并,从而减少流水线停滞。
  • 约束条件: 与第一阶段类似,确保样本分配、适配器特定词元计数和箱子容量约束。

优化效率:

  • 超时回退: 为 MILP 求解器设置超时,如果超时则回退到贪婪装箱算法 (Algorithm 1,第 2, 5, 9 行)。
  • 并行处理: 由于全局批次是独立的,可以使用多进程 (multiprocessing) 并行化批处理优化 (Algorithm 1,第 1 行),以高效调度所有训练数据。

4.2.5. 合并与验证 (Merging & Verification)

  • 合并: 经过微批次打包后,全局批次中的最后一个微批次可能未被完全填充 (underfilled),这会降低 GPU 效率并增加流水线气泡。为了缓解这个问题,调度器会进行一个贪婪合并 (greedy merge) 过程,将下一个全局批次的词元移到当前批次的最后一个微批次中(如图 12 底部所示),只要容量和气泡引理得到满足 (Algorithm 1,第 12-14 行)。
  • 验证: 执行验证步骤以确保没有违反任何约束。如果气泡条件未满足,则插入空操作微批次 (no-op microbatches) 以恢复正确性并保持流水线一致性 (Algorithm 1,第 15 行)。

4.2.6. 并行性分析器 (Parallelism Profiler)

调度器需要词元容量作为输入,而这取决于并行策略。LoRAFusion 假设有效的调度能使每个微批次中的词元数接近词元容量。由于词元容量和并行策略与调度是正交的,因此应该使用自动并行化技术 [51, 53, 106, 112] 在调度器外部进行调优。LoRAFusion 实现了一个轻量级分析器,直接在固定长度输入的不同模型并行配置下进行基准测试,收集吞吐量数据。它选择性能最佳的配置,并将其词元容量传递给数据批处理阶段,确保微批次打包与系统性能特征对齐。


5. 实验设置

5.1. 数据集

实验使用了三种常用的摘要任务公开数据集,它们具有不同的序列长度分布(如图 13 所示),这在现实条件下对批处理和调度带来了挑战:

  • XSum [61]: 英国广播公司新闻文章摘要数据集。

  • CNN/DailyMail [78] (CNNDM): CNN 和 Daily Mail 新闻文章摘要数据集。

  • WikiSum [12]: 维基百科文章摘要数据集。

    Figure 13. Distribution of sample lengths across the XSum \[61\], CNN/DailyMail \[78\], and WikiSum \[12\] datasets used for LoRA fine-tuning. 该图像是图表,展示了在XSum、CNN/DailyMail和WikiSum数据集上样本长度的分布情况。横轴为样本长度(以# Tokens表示),纵轴为密度。不同颜色的曲线代表各数据集的分布,虚线表示对应的均值。

以下是原文 Figure 13 的结果: Figure 13 展示了用于 LoRA 微调的 XSum、CNN/DailyMail 和 WikiSum 数据集中的样本长度分布。图中显示,这三个数据集的序列长度分布具有显著差异,XSum 倾向于较短的序列,而 CNN/DailyMail 和 WikiSum 具有更长的尾部和更宽的分布,这为批处理和调度带来了挑战。

对于多 LoRA 实验,所有配置都并行训练四个 LoRA 适配器,但数据集使用方式不同:

  • XSum, CNN/DailyMail (CNNDM), WikiSum 配置: 所有四个适配器都在相同的数据集上独立训练。
  • Mixed 配置: 每个适配器都在一个混合了三个数据集样本的数据集上训练。
  • Heterogeneous (Het) 配置: 四个适配器在不同的数据集上训练:每个适配器分别在 XSum、CNN/DailyMail、WikiSum 和 Mixed 数据集上训练。

5.2. 评估指标

论文的主要评估指标是吞吐量 (throughput),以每秒训练词元数 (trained tokens per second) 来衡量。

  • 概念定义: 每秒训练词元数 衡量的是系统在单位时间内处理的训练数据总量,其中“词元”是语言模型处理的基本单位(如单词、子词)。这个指标直接反映了训练效率,因为它考虑了批处理大小和序列长度,特别适用于处理序列长度变化的输入。更高的每秒训练词元数意味着更快的训练速度和更高的资源利用率。
  • 数学公式: 假设在一个时间段 Δt\Delta t 内,系统总共处理了 NN 个训练样本,每个样本包含 LiL_i 个词元。那么总处理词元数为 i=1NLi\sum_{i=1}^{N} L_iThroughput (tokens/sec)=i=1NLiΔt \text{Throughput (tokens/sec)} = \frac{\sum_{i=1}^{N} L_i}{\Delta t}
  • 符号解释:
    • i=1NLi\sum_{i=1}^{N} L_i: 在时间段 Δt\Delta t 内处理的所有训练样本的词元总数。

    • Δt\Delta t: 测量吞吐量的时间段,以秒为单位。

      论文强调选择此指标是因为它能更好地反映在序列长度变化输入下的系统效率,而非传统的每秒样本数,因为后者无法体现不同长度样本的计算量差异。论文主要关注系统性能优化,因此未报告如 BLEU、ROUGE 等模型质量指标,因为其优化是无损的,不影响模型收敛或最终质量。

5.3. 对比基线

LoRAFusion 与以下三个基线进行了比较:

  • Megatron-LM [81] with FSDP (Fully Sharded Data Parallelism): Megatron-LM 是一个领先的分布式训练框架。FSDP 是一种先进的数据并行策略,通过分片模型状态来减少内存使用。
  • Megatron-LM [81] with PP (Pipeline Parallelism): 同样基于 Megatron-LM,但采用流水线并行策略,将模型层分割到不同的 GPU 上。
  • mLoRA [98]: 最先进的多 LoRA 微调系统。
    • 再实现说明: 原论文指出,mLoRA 原始实现使用 Python RPC 进行 GPU 间通信,在配备 NVLink 的 GPU 上表现不佳。因此,为了进行公平比较,本文在自己的系统内使用高性能通信原语重新实现了 mLoRA。
    • 性能假设: 鉴于 mLoRA 没有提供独特的 Multi-LoRA CUDA 内核,本文乐观地假设其性能与朴素的单 LoRA 内核相同。
  • 未评估张量并行 (TP): 由于现有 LoRA 框架缺乏对张量并行的有效支持,TP 未被纳入评估范围。

5.4. 硬件设置

  • 主要基准测试平台:
    • NVIDIA H100 (80GB) GPU 集群:
      • 每个节点配备 8 块 H100 GPU,通过 NVLink 互联。
      • 每个节点 208 个 vCPU。
      • 节点间通过 InfiniBand [66] 进行通信。
  • 额外报告平台 (通用性演示):
    • NVIDIA L40S (48GB) GPU 服务器:
      • 每台服务器包含 4 块 L40S GPU,通过 PCIe 互联。
      • 128 个 vCPU。
  • GPU 使用数量: 大多数实验使用能容纳模型并保持良好利用率的最小 GPU 数量,通常是 1、2 或 4 块 GPU。

5.5. 软件设置

  • 深度学习框架: PyTorch 2.6

  • CUDA 工具包: CUDA Toolkit 12.4

  • 内核开发框架: Triton 3.2.0

  • 分布式训练框架: Megatron-Core 0.11.0

  • 模型架构和 LoRA 适配: 集成了 Hugging Face Transformers [93] 和 PEFT Library [55]。


6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 端到端结果

H100 GPU 上的加速 (Figure 14):

Figure 14. End-to-end training throughput (tokens/sec) of training 4 LoRA adapters on 1, 2, and \(4 \\mathrm { H } 1 0 0\) GPUs. The first four b He on different datasets. 该图像是图表,展示了在不同数据集上,使用 1、2 和 4 个 H100 GPU 训练 4 个 LoRA 适配器的端到端训练吞吐量(tokens/sec)。图中比较了不同方法的吞吐量,包括 mLoRA 和 LoRAFusion。

以下是原文 Figure 14 的结果: Figure 14 展示了在 1、2 和 4 块 H100 GPU 上训练 4 个 LoRA 适配器的端到端训练吞吐量(tokens/sec)。LoRAFusion 始终优于 Megatron-LM 和 mLoRA,在 LLaMa-3.1-8B、Qwen-2.5-32B 和 LLaMa-3.1-70B 模型上,实现了显著的加速。特别是对于 WikiSum 数据集,LoRAFusion 的表现尤为突出,因为其能够处理高度可变的样本长度。

  • LoRAFusion 在所有模型和数据集上都持续优于所有基线方法,加速比范围为 1.19×1.19 \times1.96×1.96 \times
  • LLaMa-3.1-8B (单块 H100 GPU): LoRAFusion 实现了平均 1.26×1.26 \times(最高 1.43×1.43 \times)的加速。这主要得益于 FusedLoRA 内核减少了内存流量。由于单 GPU 设置不存在负载不平衡问题,此处的改进直接反映了内核层面的收益。在 WikiSum 数据集上实现了高加速,因为该数据集样本长度变化大,而 LoRAFusion 能够稳定地进行打包,而基线方法则面临内存不足 (OOM) 或性能下降。
  • Qwen-2.5-32B 和 LLaMa-3.1-70B (分布式训练): LoRAFusion 分别实现了平均 1.42×1.42 \times1.64×1.64 \times(最高 1.64×1.64 \times1.81×1.81 \times)的加速。模型越大,从改进的调度中获益越多,因为流水线停滞和负载不平衡在更高的并行度下变得更明显。
  • 异构 (Het) 设置: 在最具挑战性的异构设置中(每个适配器使用不同的数据集),LoRAFusion 仍然表现出强大的性能,突显了其鲁棒性。

L40S GPU 上的加速 (Figure 15):

Figure 15. End-to-end training throughput (tokens/sec) of training 4 LoRA adapters on 1 and 4 L40S GPUs. 该图像是一个柱状图,展示了在 Llama-3.1-8B 和 Qwen-2.5-32B 模型下,使用不同方法(如 LoRAFusion、mLoRA 等)进行混合和异构训练的吞吐量(tokens/sec)。图中显示了 LoRAFusion 在两种模型下的显著性能提升。

以下是原文 Figure 15 的结果: Figure 15 展示了在 1 块和 4 块 L40S GPU 上训练 4 个 LoRA 适配器的端到端训练吞吐量(tokens/sec),涵盖了 LLaMa-3.1-8B 和 Qwen-2.5-32B 模型。LoRAFusion 在 L40S 平台上也展现了显著的加速,在 LLaMa-3.1-8B 上平均加速 1.19×1.19 \times,在 Qwen-2.5-32B 上平均加速 1.91×1.91 \times,验证了其在不同硬件平台上的通用性。

  • LoRAFusion 在 NVIDIA L40S GPU 上也取得了 1.19×1.91×1.19 \times - 1.91 \times 的平均加速。
  • 对于 LLaMa-3.1-8B,由于单块 L40S GPU 内存容量有限,限制了批处理大小并影响了内核融合的有效性,因此收益较小。然而,即使在这种约束下,LoRAFusion 仍保持了持续的改进,这表明其在不同模型规模和硬件平台上的通用性。

6.1.2. 可扩展性研究

H100 GPU 上的可扩展性 (Figure 16):

Figure 16. Scalability of LoRAFusion across 4, 8, and 16 H100 GPUs when training 4 LoRA adapters simultaneously. DP scaling means the more GPUs are used to increase the DP degree for the same job, while Job scaling means different LoRA fine-tuning jobs are scheduled to utilize more GPUs. Global batch sizes are scaled proportionally with GPU count to ensure fair comparison. 该图像是一个图表,展示了在同时训练4个LoRA适配器时,LoRAFusion与Megatron-LM的性能对比。图中列出了不同GPU数量(4、8、16)下的吞吐量(K tokens/s)以及相应的缩放因子,显示了LoRAFusion在DP和Job缩放中的优越性能。

以下是原文 Figure 16 的结果: Figure 16 展示了 LoRAFusion 在训练 4 个 LoRA 适配器时,在 4、8 和 16 块 H100 GPU 上的可扩展性。图中比较了两种扩展策略:数据并行 (DP) 扩展(每个任务使用更多 GPU)和任务级 (Job) 扩展(调度更多并发任务)。LoRAFusion 在 Job 扩展模式下显示出更高的吞吐量,同时在 DP 扩展下也优于基线。

  • 研究了 4、8 和 16 块 H100 GPU 上的两种扩展策略:DP 扩展(每个任务使用更多 GPU)和任务级 (Job) 扩展(运行更多并发任务)。全局批处理大小与 GPU 数量成比例缩放。
  • 主要结论:
    1. 任务级扩展持续优于 DP 扩展: Job 扩展实现了更好的负载均衡,在 8 和 16 块 GPU 上分别实现了 1.18×1.18 \times1.25×1.25 \times 的更高吞吐量。这表明为每个任务分配更少的 GPU 而使用更多 GPU 来运行更多独立任务通常能带来更好的效率,因为它减少了 GPU 间通信和同步开销。
    2. LoRAFusion 与 DP 扩展和多节点微调完全兼容: 即使在 DP 扩展下,LoRAFusion 仍然表现出强大的性能,比 Megatron-LM 平均加速 1.78×1.78 \times,比 mLoRA 平均加速 1.50×1.50 \times。这证明了 LoRAFusion 在不同并行配置下的通用性和有效性。

6.1.3. FusedLoRA 内核的有效性

内核性能 (Figure 17):

Figure 17. Performance of FusedLoRA kernel in forward and backward passes. 该图像是图表,展示了不同模型(Torch LoRA、FusedLoRA、FusedMultiLoRA)在不同 Token 数量下的规范化吞吐量。图中显示了在 N=K=4096、N=K=5120 和 N=K=8192 的条件下,各模型的性能变化趋势,FusedLoRA 在大多数情况下表现优越。

以下是原文 Figure 17 的结果: Figure 17 比较了 FusedLoRA 和 FusedMultiLoRA 内核与标准 PyTorch LoRA 实现的吞吐量,涵盖前向和反向传播,并在不同 Token 数量和模型配置下进行了测试。结果显示 FusedLoRA 和 FusedMultiLoRA 均显著优于基线,特别是在 N=K=4096 和 N=K=8192 的配置下,它们的吞吐量表现出显著优势。

  • FusedLoRA 内核比标准的 Torch LoRA 实现平均加速 1.27×1.27 \times(最高 1.39×1.39 \times)。
  • FusedMultiLoRA 内核平均加速 1.17×1.17 \times(最高 1.24×1.24 \times)。
  • 在前向传播中,FusedMultiLoRA 性能与 FusedLoRA 相似,因为大部分计算是共享的。
  • 在反向传播中,FusedMultiLoRA 存在轻微开销,这源于跨适配器累积梯度和额外的逐元素操作。
  • 尽管存在这些开销,两个内核在不同词元大小和模型配置下都持续优于基线。

层级性能 (Figure 18):

Figure 18. Performance of FusedLoRA kernel in decoder layers of different models. 该图像是一个图表,展示了不同模型在不同批量大小下的标准化吞吐量。图表中包括了三种方法:Torch LoRA、FusedLoRA 和 FusedMultiLoRA,分别针对 Llama-3.1-8B、Qwen2.5-32B 和 Llama-3.1-70B 模型进行比较。

以下是原文 Figure 18 的结果: Figure 18 比较了 FusedLoRA 和 FusedMultiLoRA 内核在不同模型解码器层中的性能加速。图中显示,FusedLoRA 在 Llama-3.1-8B、Qwen2.5-32B 和 Llama-3.1-70B 模型上平均加速 1.21×1.21 \times,而 FusedMultiLoRA 平均加速 1.13×1.13 \times,证明了融合内核在实际模型架构中的有效性。

  • FusedLoRA 在模型不同线性层中平均加速 1.21×1.21 \times(最高 1.30×1.30 \times)。
  • FusedMultiLoRA 平均加速 1.13×1.13 \times(最高 1.17×1.17 \times)。
  • 这些结果基于包含四个适配器的微批次。在实际微调工作负载中,每个微批次通常只包含一两个适配器,这使得 FusedMultiLoRA 的性能接近 FusedLoRA

内存流量减少 (Figure 19):

该图像是一个示意图,展示了不同 LoRA 方法在 DRAM 读取/写入性能上的比较。数据显示,在不同的维度配置下,Fused LoRA 和 Fused MultiLoRA 相较于传统的 Torch LoRA 具有更低的内存使用量。具体表现为在矩阵尺寸 \(8192 \\times 4096 \\times 4096\) 时,Torch LoRA 为 1.00 倍,而 Fused LoRA 为 0.63 倍。 该图像是一个示意图,展示了不同 LoRA 方法在 DRAM 读取/写入性能上的比较。数据显示,在不同的维度配置下,Fused LoRA 和 Fused MultiLoRA 相较于传统的 Torch LoRA 具有更低的内存使用量。具体表现为在矩阵尺寸 8192×4096×40968192 \times 4096 \times 4096 时,Torch LoRA 为 1.00 倍,而 Fused LoRA 为 0.63 倍。

以下是原文 Figure 19 的结果: Figure 19 展示了使用 NVIDIA Nsight Compute (NCU) 对不同内核(Torch LoRA、FusedLoRA、FusedMultiLoRA)的 GPU DRAM 内存流量进行的比较。结果表明,FusedLoRA 和 FusedMultiLoRA 在各种 GEMM 形状下都显著减少了内存读写,例如,对于 8192×4096×40968192 \times 4096 \times 4096 的形状,总 DRAM 流量减少到 0.63×0.63 \times,验证了融合设计有效减少了冗余内存访问。

  • FusedLoRAFusedMultiLoRA 持续减少了与 Torch LoRA 相比的内存使用。
  • 例如,对于 8192×4096×40968192 \times 4096 \times 4096 的 GEMM 形状,总 DRAM 流量减少到 0.63×0.63 \times
  • 在所有设置下,内存流量减少了 34%37%34\% - 37\%,这证实了融合设计有效地减少了冗余内存访问。
  • 性能洞察: 融合内核减少了大型激活张量的冗余内存访问,这在内存带宽远低于计算浮点运算的硬件上尤为重要。随着现代加速器计算能力增长快于内存带宽,融合内核的收益预计在未来系统中会更大。

6.1.4. 任务级调度的有效性

流水线气泡减少 (Figure 20):

Figure 19. GPU DRAM memory traffic comparison between different kernels from NVIDIA Nsight Compute (NCU). 该图像是图表,展示了不同方法下的管道气泡比率。Megatron-LM表现出48.79%,而mLoRA则为34.11%。在使用一个适配器的LoRAFusion中,管道气泡比率为44.17%;使用两个适配器的LoRAFusion为15.00%;三个适配器为12.23%;而四个适配器的比率则降至11.09%。

以下是原文 Figure 19 的结果: Figure 19(原文 Figure 20)展示了不同方法下的流水线气泡比率。Megatron-LM 的气泡比率最高,为 48.79%。mLoRA 略有改善,为 34.11%。LoRAFusion 的气泡比率随着并行适配器数量的增加而显著下降:从单个适配器的 44.17% 降至四个适配器时的 11.09%。这表明 LoRAFusion 的调度策略在减少流水线空闲时间方面非常有效。

  • 单适配器: 当只有一个适配器时,气泡比率高达 44.17%44.17\%,接近 Megatron-LM 的 48.79%48.79\%。这表明当只有一个数据集可用时,分组是无效的,突显了多 LoRA 对于改进调度灵活性的重要性。
  • 多适配器: 随着更多适配器一起训练,气泡比率稳步下降:2 个适配器时为 15.00%15.00\%,3 个适配器时为 12.23%12.23\%,4 个适配器时为 11.09%11.09\%
  • 对比 mLoRA: mLoRA 仅达到 34.11%34.11\% 的气泡比率,这证实了 LoRAFusion 的分组和批处理显著减少了流水线空闲时间。
  • 剩余气泡: 即使有四个适配器,气泡比率仍为 11.09%11.09\%。这主要是由于流水线阶段执行时间不均匀,最后一个阶段通常需要更长的时间,因为它处理额外的线性层和交叉熵损失。

调优时间 (Figure 21):

Figure 21. Tuning and computation time vs. number of samples for 4-stage pipeline with 4 adapters. 该图像是图表,展示了在具有4个适配器的4阶段流水线中,调优时间与样本数量之间的关系。随着样本数量的增加,计算时间呈现出明显递增的趋势,而调优时间则相对平稳。

以下是原文 Figure 21 的结果: Figure 21 描绘了在 4 阶段流水线、4 个适配器的情况下,调优时间与训练样本数量的关系。随着训练样本数量从 640 增加到 25600,计算时间几乎线性增长,而调度时间也呈线性增长,但增长斜率远小于计算时间,表明调度开销相对较小且可控。

  • 在 64 个 vCPU 和 4 块 H100 GPU 上,对于 4 阶段流水线和 4 个适配器,调优和计算时间随训练样本数量的增长而变化。
  • 调度时间: 调度时间几乎呈线性增长,从 640 个样本时的 15.74 秒到 25600 个样本时的 102.12 秒。这表明调度器具有线性可扩展性。
  • 计算时间: 计算时间也几乎呈线性增长,但斜率远大于调度时间。
  • 调度开销: 调度开销可以忽略不计,原因有三:
    1. CPU 上的调度与 GPU 上前一个全局批次的训练并行运行,调度器的延迟可以被这种重叠完全隐藏。
    2. 性能收益在 4 个适配器时趋于饱和,允许在实际部署中以少量固定数量的适配器运行。
    3. 对 MILP 求解器设置了超时,如果求解时间过长则回退到贪婪算法,确保调度开销始终在可控范围内。

合并与贪婪回退的有效性:

  • 合并 (merging) 过程使吞吐量提高了 4.34%4.34\%
  • 两阶段 MILP 优化比纯贪婪装箱算法额外提供了 3.82%3.82\% 的改进。
  • MILP 求解器路径在 10 秒超时下被 77.4%77.4\% 的全局批次选择,表明其在减少未填充微批次中的词元数方面是有效的。
  • 这些适度的改进表明,大多数微批次已经打包良好,算法主要优化每个全局批次中的最后一个微批次。由于调度开销通过并行 GPU 执行而被隐藏,这些优化在不引入额外开销的情况下将性能推向硬件极限。

6.1.5. 加速分解

LoRAFusion 在 LLaMa-3.1-70B (4 块 GPU) 上的加速分解 (Figure 22):

Figure 22. Speedup breakdown of LoRAFusion on LLaMa3.1-70B with 4 GPUs. 该图像是图表,展示了在使用4个GPU时,LoRAFusion在LLaMa3.1-70B上的相对加速效果。各项加速效果与1F1B PP的基准1.00x进行对比,其中平衡的Multi-LoRA ZeRO Bubble PP + FusedMultiLoRA达到最高的2.05x加速效果。

以下是原文 Figure 22 的结果: Figure 22 展示了 LoRAFusion 在 LLaMa-3.1-70B 模型(使用 4 块 GPU)上的加速分解。基线是 Megatron-LM 的 1F1B PP,其吞吐量为 1.00x。单独添加 FusedLoRA 带来 1.13x 加速。替换为 Multi-LoRA zero-bubble PP 将加速提升到 1.50x。在此基础上再添加 FusedMultiLoRA 进一步提升至 1.72x。仅应用调度器(Adaptive Scheduling)即可实现 1.57x 加速。最终,结合自适应调度和融合内核,LoRAFusion 实现了 2.05x 的最高加速。

  • 基线 (1F1B PP): Megatron-LM 中使用的单前向单反向流水线并行。
  • + FusedLoRA: 仅添加 FusedLoRA 内核,带来了 1.13×1.13 \times 的加速。这表明内核融合本身就能带来收益,但由于负载不平衡和次优的词元形状,其效率受到限制。
  • + Multi-LoRA zero-bubble PP: 将 1F1B 替换为 Multi-LoRA 零气泡流水线并行(通过来自独立适配器的更多微批次消除流水线停滞),吞吐量提高到 1.50×1.50 \times
  • + FusedMultiLoRA (在此基础上): 在 Multi-LoRA zero-bubble PP 的基础上添加 FusedMultiLoRA 内核,进一步将加速提升到 1.72×1.72 \times,因为它支持多适配器微批次并减少了冗余内存访问。
  • + Adaptive Scheduling (独立): 仅应用调度器来重新平衡微批次间的词元分布,即使不使用融合内核,性能也提高到 1.57×1.57 \times,因为它显著减少了负载不平衡。
  • LoRAFusion (自适应调度 + 融合内核): 结合自适应调度和融合内核,实现了最高的 2.05×2.05 \times 加速。这表明联合优化内核效率、并行性和工作负载平衡的重要性。

与 mLoRA 相比的加速原因:

  • 内核融合: 我们的内核融合带来了 1.15×1.15 \times 的加速(比较 Figure 22 中的第 3 栏和第 4 栏)。当序列长度规则且与内核性能匹配时,收益更大。
  • 自适应批处理: 我们的自适应批处理缓解了负载不平衡,带来了 1.19×1.19 \times 的加速(比较 Figure 22 中的第 4 栏和第 6 栏)。
  • 这些改进得到了微基准测试的进一步支持:Figure 17 显示内核性能平均加速 1.17×1.17 \times(最高 1.24×1.24 \times),Figure 20 显示流水线气泡减少了 23.02%23.02\%

6.2. 数据呈现 (表格)

以下是原文 Table 1 的结果:

符号 描述
r LoRA 秩
m 词元数(批处理大小 X 序列长度)
k,n 权重矩阵的输入/输出维度
W 基础模型权重。大小: (k, n)。
A 第一个 LoRA 权重。大小: (k, r)
B 第二个 LoRA 权重。大小: (r, n)
X 输入张量。大小: (m, k)
X^\widehat{X} Dropout 后的输入张量。大小: (m, k)
Y 输出张量。大小: (m, n)

7. 总结与思考

7.1. 结论总结

本文识别并解决了 LLM LoRA 微调中两个关键的性能瓶颈:LoRA 模块中冗余的内存访问和多任务并发 LoRA 微调中错失的优化机会。LoRAFusion 作为解决方案,引入了针对 LoRA 内核的新型水平融合技术,将内存流量减少了高达 37%37\%。同时,其互补的任务级调度策略将 GPU 利用率从 65%65\% 提高到 89%89\%。这些优化相结合,使得 LoRAFusion 相比最先进的系统实现了高达 1.96×1.96 \times 的加速(平均 1.47×1.47 \times),并在各种模型和数据集上表现出色。FusedLoRA 内核作为即插即用组件,可以直接集成到现有系统中,提供即时收益。LoRAFusion 的贡献在于通过系统级的创新,显著提升了 LLM LoRA 微调的可访问性和效率。

7.2. 局限性与未来工作

论文作者指出了以下局限性并提出了未来研究方向:

  • LoRA 变体的通用性: 论文的内核融合设计可以扩展到 DoRA [52] 和 VeRA [37] 等其他流行的 LoRA 变体。这些变体通常在核心 LoRA 计算前后添加预处理或后处理函数。LoRAFusion 的优化与这些修改是正交的,用户可以通过定义 prologue/epilogue 函数来扩展其内核。然而,手动扩展效率不高。
    • 未来工作: 计划将融合模式集成到编译器框架中,利用 torch.compile 添加编译器注释作为提示,指导 LoRA 模式的融合过程。这将实现 LoRA 变体优化的自动化,无需手动内核开发。
  • 量化方法的通用性: LoRAFusion 提出的内核可以直接应用于 4 比特 QLoRA [14]。当前的 QLoRA 实现通常在 LoRA 计算前将 4 比特权重反量化 (dequantize) 为半精度,这使得 LoRAFusion 的内核无需修改即可工作。虽然反量化也可以与 LoRA 路径融合,但近期研究表明,对于大量词元,两步法通常性能更优 [17]。
  • 流水线气泡的根本限制: 即使 LoRAFusion 显著减少了流水线气泡(降至 11.09%11.09\%),仍有部分气泡存在。这主要是由于流水线阶段执行时间不均匀,特别是最后一个阶段通常需要更长的时间来处理额外的线性层和交叉熵损失。这部分限制是调度器无法解决的,超出了本文的范围。

7.3. 个人启发与批判

7.3.1. 个人启发

  • 系统级优化的重要性: 论文强调了系统级优化(从低层内核到高层调度)对于提升 LLM 训练效率的巨大潜力。仅仅依靠算法改进或单一层面的优化是不够的,需要多层级的协同设计。
  • 内存带宽是关键瓶颈: 论文通过详细的剖析,明确指出 LoRA 微调的主要瓶颈并非计算量(FLOPs),而是对大型激活张量的冗余内存访问。这提醒我们在设计和优化深度学习系统时,内存带宽往往比浮点运算能力更早成为瓶颈,特别是在参数高效模型中。
  • 变长序列处理的挑战与机遇: 真实世界的训练数据往往包含变长序列,这导致批处理和分布式训练中的负载不平衡。LoRAFusion 的调度器通过智能分组和装箱策略,将这种挑战转化为优化机会,通过平衡负载和减少流水线气泡来提高整体吞吐量。
  • Multi-LoRA 的多重收益: Multi-LoRA 不仅可以减少内存占用,其更大的价值在于为调度器提供了更多的灵活性和并行度,从而能够有效缓解分布式训练中的通信开销和负载不平衡。
  • 编译器优化的前景: 论文提到的将融合模式集成到编译器框架中的未来工作,预示着自动化和智能化的编译器技术将在未来深度学习系统优化中扮演越来越重要的角色,降低开发者手动优化内核的门槛。

7.3.2. 批判与潜在问题

  • 调度器复杂性与开销的平衡: 尽管论文指出调度器开销可控且可被 GPU 训练隐藏,但 MILP 求解器本身是计算密集型的。对于超大规模、极度动态(任务频繁进出)或任务数量极多的场景,调度器在 CPU 上的预处理时间仍可能成为瓶颈,即使有超时和贪婪回退。未来可以探索更轻量级或增量的调度策略。
  • 泛化到其他任务类型: 论文主要在摘要任务上进行评估。虽然摘要任务具有可变序列长度的特性,但 LLM 的其他任务(如代码生成、对话)可能具有不同的特征,例如更长的生成序列、更复杂的依赖关系或不同的计算-内存访问模式。LoRAFusion 在这些任务上的性能和调度策略的有效性可能需要进一步验证。
  • 对并行策略的依赖性: 调度器需要预先知道最佳的并行策略(通过轻量级分析器)。虽然这是合理的工程选择,但如果最佳并行策略本身非常动态或难以确定,可能会影响调度器的整体效能。未来的研究可以探索将并行策略搜索与调度本身更紧密地集成。
  • 特定硬件的优化: LoRAFusion 的内核优化基于 Triton,虽然通用性较好,但其最佳性能依赖于硬件特定的调优配置 (tune_kernels.py)。这可能意味着在新的硬件平台部署时需要额外的调优工作,而非完全的即插即用。
  • 残余流水线气泡的根源: 尽管 LoRAFusion 显著减少了流水线气泡,但仍有 11.09%11.09\% 的气泡是由于模型架构本身在不同阶段计算量不均衡造成的。这暗示了在模型设计或流水线切分策略上可能仍有改进空间,以实现更均衡的负载分布,但这超出了系统优化的范畴,更多属于算法或架构层面的挑战。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。