大语言模型(LLMs)推理过程中的批处理优化面临显著挑战,这主要源于其推理过程的迭代特性。核心问题在于批处理中的各个请求完成时间存在差异,这导致资源释放和新请求整合的复杂性显著提高,特别是在处理不同完成阶段的请求时。当批处理中序列的生成长度差异较大时,GPU资源利用率往往不够理想。如下图所示,序列1、3和4中,结束符号(EOS)之后的白色区块直观地展示了GPU资源的闲置状态,凸显了当前批处理方法的效率损失。
ORCA系统创新性地提出了持续批处理概念,通过引入迭代级调度和选择性批处理机制,有效解决了大语言模型批处理中的关键技术挑战。
特性1 --- 迭代级调度
下图详细展示了基于迭代级调度的ORCA系统架构及其工作流程。ORCA系统对外提供标准化接口(支持HTTPS或gRPC协议)用于处理推理请求。系统将接收到的请求存储于请求池中,由调度器进行统一监控和管理。调度器的核心职责包括:从请求池中筛选请求、调度执行引擎执行模型迭代、接收引擎输出的标记序列,并据此更新请求池状态。执行引擎负责进行张量运算,支持跨多GPU和多机器的并行处理能力。与传统批处理方法不同,ORCA的调度器能够在每次迭代中动态调整处理请求的组合,提供了更灵活的请求处理机制。当某个请求处理完成后,系统将其从请求池中移除,并通过接口返回响应结果。这种方法突破了传统批处理必须等待所有请求完成全部迭代的限制。
特性2 --- 选择性批处理
选择性批处理技术通过在张量运算中实现差异化的批处理策略,显著提升了模型处理的灵活性。这种创新方法特别适用于处理形状不规则的输入张量场景,这类张量通常难以直接合并为统一的大规模张量进行批处理。在非注意力机制的矩阵乘法运算和层归一化等操作中,选择性批处理采用将不规则张量重构为无批处理维度的二维结构的方式,实现统一的大规模矩阵计算。这种优化策略的可行性基于这些操作的逐标记处理特性,使得来自不同请求的标记可以等同处理。但是对于需要请求专属键值对的注意力机制操作,系统采用分离式批处理策略,确保每个请求独立处理。注意力操作完成后,输出结果会重新合并为统一张量用于后续处理。ORCA通过K/V管理器持久化存储先前注意力操作产生的键值对,使增量请求能够在多轮迭代中复用这些信息。
分布式架构设计
ORCA系统在Transformer模型并行化方面融合了两种关键技术:层内并行与层间并行。层内并行技术(这也是FasterTransformer采用的方法)主要针对矩阵乘法运算(包括线性变换和注意力机制)及其参数在多GPU间进行分配。而层间并行则致力于在多个GPU之间分配Transformer层。ORCA采用均衡分配策略,确保每个GPU负责处理相同数量的Transformer层。例如,对于一个4层GPT模型,可以先划分为2个层间并行分区,每个分区再细分为3个层内并行分区,最终形成6个GPU的协同计算架构。
FasterTransformer是NVIDIA开发的开源优化库,专注于加速Transformer模型,支持BERT、GPT-2、GPT-J和T5等主流架构。目前,FasterTransformer的技术演进已转移至TensorRT-LLM平台。
ORCA执行引擎采用分布式执行策略,将工作进程分配到不同物理机器上,每个工作进程负责管理一个或多个控制GPU的CPU线程。在具体实现中,每个工作进程专注于处理模型的特定层间分区,而线程数量则由层内并行度决定。当系统为一批请求调度模型迭代时,引擎主控程序首先向第一个工作进程(Worker1)传递必要的计算信息,包括标记序列和控制指令。Worker1随后将这些信息分发给其管理的GPU控制线程,触发相应的GPU计算内核。同时,Worker1采用非阻塞方式将控制指令转发给下一个工作进程(Worker2)。最后一个工作进程在确保其GPU任务全部完成后,负责汇总输出标记并将结果返回给引擎主控程序。这种架构设计实现了跨设备、跨机器的高效分布式计算。
调度算法设计
ORCA调度器采用基于到达时间的优先级机制,在最大批处理大小约束下动态选择请求进行处理。系统在首次调度请求时会预分配存储键值对所需的内存空间。调度过程中
Select
函数确保严格遵循请求的时间优先级进行批次选择。调度器通过评估最大标记数需求,动态验证新请求的内存资源可用性。在确认内存充足的情况下,系统保证为请求的完整生命周期预留必要的键值对缓存空间。整个调度机制采用并发处理模式,支持请求的动态添加和完成请求的及时移除。
ORCA调度器通过任务流水线化技术优化工作进程执行效率。系统不采用传统的串行批次处理模式,而是维持并发运行批次数与工作进程数的动态平衡。这种策略确保了工作进程的持续高效运转。例如在三工作进程配置下,最大批处理大小为二的场景中,调度器可以同时处理多个批次(如AB、CD、EF)。在调度第一个批次(AB)后,系统立即安排后续两个批次(CD和EF)的处理,无需等待AB批次的完成。当AB批次处理完成后,如果这些请求仍具有时间优先性,则会被重新纳入调度流程。这种机制显著提升了系统的资源利用率和整体吞吐性能。
性能评估
得益于创新的迭代级调度和选择性批处理机制,ORCA系统在保持相同延迟水平的前提下,实现了相比FasterTransformer高达37倍的吞吐量提升,展现了显著的性能优势。
论文:
https://avoid.overfit.cn/post/b9a179b6cee448dba7ace15ac95cfe59
作者:Don Moon