ORCA：基于持續(xù)批處理的LLM推理性能優(yōu)化技術(shù)詳解

作者：佚名 2024-12-09 13:40:26

大語(yǔ)言模型(LLMs)推理過(guò)程中的批處理優(yōu)化面臨顯著挑戰(zhàn)，這主要源于其推理過(guò)程的迭代特性。核心問(wèn)題在于批處理中的各個(gè)請(qǐng)求完成時(shí)間存在差異，這導(dǎo)致資源釋放和新請(qǐng)求整合的復(fù)雜性顯著提高，特別是在處理不同完成階段的請(qǐng)求時(shí)。當(dāng)批處理中序列的生成長(zhǎng)度差異較大時(shí)，GPU資源利用率往往不夠理想。如下圖所示，序列1、3和4中，結(jié)束符號(hào)（EOS）之后的白色區(qū)塊直觀地展示了GPU資源的閑置狀態(tài)，凸顯了當(dāng)前批處理方法的效率損失。

ORCA系統(tǒng)創(chuàng)新性地提出了持續(xù)批處理概念，通過(guò)引入迭代級(jí)調(diào)度和選擇性批處理機(jī)制，有效解決了大語(yǔ)言模型批處理中的關(guān)鍵技術(shù)挑戰(zhàn)。

特性1 — 迭代級(jí)調(diào)度

下圖詳細(xì)展示了基于迭代級(jí)調(diào)度的ORCA系統(tǒng)架構(gòu)及其工作流程。ORCA系統(tǒng)對(duì)外提供標(biāo)準(zhǔn)化接口（支持HTTPS或gRPC協(xié)議）用于處理推理請(qǐng)求。系統(tǒng)將接收到的請(qǐng)求存儲(chǔ)于請(qǐng)求池中，由調(diào)度器進(jìn)行統(tǒng)一監(jiān)控和管理。調(diào)度器的核心職責(zé)包括：從請(qǐng)求池中篩選請(qǐng)求、調(diào)度執(zhí)行引擎執(zhí)行模型迭代、接收引擎輸出的標(biāo)記序列，并據(jù)此更新請(qǐng)求池狀態(tài)。執(zhí)行引擎負(fù)責(zé)進(jìn)行張量運(yùn)算，支持跨多GPU和多機(jī)器的并行處理能力。與傳統(tǒng)批處理方法不同，ORCA的調(diào)度器能夠在每次迭代中動(dòng)態(tài)調(diào)整處理請(qǐng)求的組合，提供了更靈活的請(qǐng)求處理機(jī)制。當(dāng)某個(gè)請(qǐng)求處理完成后，系統(tǒng)將其從請(qǐng)求池中移除，并通過(guò)接口返回響應(yīng)結(jié)果。這種方法突破了傳統(tǒng)批處理必須等待所有請(qǐng)求完成全部迭代的限制。

特性2 — 選擇性批處理

選擇性批處理技術(shù)通過(guò)在張量運(yùn)算中實(shí)現(xiàn)差異化的批處理策略，顯著提升了模型處理的靈活性。這種創(chuàng)新方法特別適用于處理形狀不規(guī)則的輸入張量場(chǎng)景，這類張量通常難以直接合并為統(tǒng)一的大規(guī)模張量進(jìn)行批處理。在非注意力機(jī)制的矩陣乘法運(yùn)算和層歸一化等操作中，選擇性批處理采用將不規(guī)則張量重構(gòu)為無(wú)批處理維度的二維結(jié)構(gòu)的方式，實(shí)現(xiàn)統(tǒng)一的大規(guī)模矩陣計(jì)算。這種優(yōu)化策略的可行性基于這些操作的逐標(biāo)記處理特性，使得來(lái)自不同請(qǐng)求的標(biāo)記可以等同處理。但是對(duì)于需要請(qǐng)求專屬鍵值對(duì)的注意力機(jī)制操作，系統(tǒng)采用分離式批處理策略，確保每個(gè)請(qǐng)求獨(dú)立處理。注意力操作完成后，輸出結(jié)果會(huì)重新合并為統(tǒng)一張量用于后續(xù)處理。ORCA通過(guò)K/V管理器持久化存儲(chǔ)先前注意力操作產(chǎn)生的鍵值對(duì)，使增量請(qǐng)求能夠在多輪迭代中復(fù)用這些信息。

分布式架構(gòu)設(shè)計(jì)

ORCA系統(tǒng)在Transformer模型并行化方面融合了兩種關(guān)鍵技術(shù)：層內(nèi)并行與層間并行。層內(nèi)并行技術(shù)（這也是FasterTransformer采用的方法）主要針對(duì)矩陣乘法運(yùn)算（包括線性變換和注意力機(jī)制）及其參數(shù)在多GPU間進(jìn)行分配。而層間并行則致力于在多個(gè)GPU之間分配Transformer層。ORCA采用均衡分配策略，確保每個(gè)GPU負(fù)責(zé)處理相同數(shù)量的Transformer層。例如，對(duì)于一個(gè)4層GPT模型，可以先劃分為2個(gè)層間并行分區(qū)，每個(gè)分區(qū)再細(xì)分為3個(gè)層內(nèi)并行分區(qū)，最終形成6個(gè)GPU的協(xié)同計(jì)算架構(gòu)。

FasterTransformer是NVIDIA開(kāi)發(fā)的開(kāi)源優(yōu)化庫(kù)，專注于加速Transformer模型，支持BERT、GPT-2、GPT-J和T5等主流架構(gòu)。目前，F(xiàn)asterTransformer的技術(shù)演進(jìn)已轉(zhuǎn)移至TensorRT-LLM平臺(tái)。

ORCA執(zhí)行引擎采用分布式執(zhí)行策略，將工作進(jìn)程分配到不同物理機(jī)器上，每個(gè)工作進(jìn)程負(fù)責(zé)管理一個(gè)或多個(gè)控制GPU的CPU線程。在具體實(shí)現(xiàn)中，每個(gè)工作進(jìn)程專注于處理模型的特定層間分區(qū)，而線程數(shù)量則由層內(nèi)并行度決定。當(dāng)系統(tǒng)為一批請(qǐng)求調(diào)度模型迭代時(shí)，引擎主控程序首先向第一個(gè)工作進(jìn)程（Worker1）傳遞必要的計(jì)算信息，包括標(biāo)記序列和控制指令。Worker1隨后將這些信息分發(fā)給其管理的GPU控制線程，觸發(fā)相應(yīng)的GPU計(jì)算內(nèi)核。同時(shí)，Worker1采用非阻塞方式將控制指令轉(zhuǎn)發(fā)給下一個(gè)工作進(jìn)程（Worker2）。最后一個(gè)工作進(jìn)程在確保其GPU任務(wù)全部完成后，負(fù)責(zé)匯總輸出標(biāo)記并將結(jié)果返回給引擎主控程序。這種架構(gòu)設(shè)計(jì)實(shí)現(xiàn)了跨設(shè)備、跨機(jī)器的高效分布式計(jì)算。

調(diào)度算法設(shè)計(jì)

ORCA調(diào)度器采用基于到達(dá)時(shí)間的優(yōu)先級(jí)機(jī)制，在最大批處理大小約束下動(dòng)態(tài)選擇請(qǐng)求進(jìn)行處理。系統(tǒng)在首次調(diào)度請(qǐng)求時(shí)會(huì)預(yù)分配存儲(chǔ)鍵值對(duì)所需的內(nèi)存空間。調(diào)度過(guò)程中Select函數(shù)確保嚴(yán)格遵循請(qǐng)求的時(shí)間優(yōu)先級(jí)進(jìn)行批次選擇。調(diào)度器通過(guò)評(píng)估最大標(biāo)記數(shù)需求，動(dòng)態(tài)驗(yàn)證新請(qǐng)求的內(nèi)存資源可用性。在確認(rèn)內(nèi)存充足的情況下，系統(tǒng)保證為請(qǐng)求的完整生命周期預(yù)留必要的鍵值對(duì)緩存空間。整個(gè)調(diào)度機(jī)制采用并發(fā)處理模式，支持請(qǐng)求的動(dòng)態(tài)添加和完成請(qǐng)求的及時(shí)移除。

ORCA調(diào)度器通過(guò)任務(wù)流水線化技術(shù)優(yōu)化工作進(jìn)程執(zhí)行效率。系統(tǒng)不采用傳統(tǒng)的串行批次處理模式，而是維持并發(fā)運(yùn)行批次數(shù)與工作進(jìn)程數(shù)的動(dòng)態(tài)平衡。這種策略確保了工作進(jìn)程的持續(xù)高效運(yùn)轉(zhuǎn)。例如在三工作進(jìn)程配置下，最大批處理大小為二的場(chǎng)景中，調(diào)度器可以同時(shí)處理多個(gè)批次（如AB、CD、EF）。在調(diào)度第一個(gè)批次（AB）后，系統(tǒng)立即安排后續(xù)兩個(gè)批次（CD和EF）的處理，無(wú)需等待AB批次的完成。當(dāng)AB批次處理完成后，如果這些請(qǐng)求仍具有時(shí)間優(yōu)先性，則會(huì)被重新納入調(diào)度流程。這種機(jī)制顯著提升了系統(tǒng)的資源利用率和整體吞吐性能。