如何在TorchServe上提供LLMs的分布式推理

這里是Hamid，我來(lái)自PyTorch合作伙伴工程部。我將跟隨Mark的討論，講解如何在TorchServe上提供LLMs的分布式推理和其他功能。首先，為什么需要分布式推理呢？簡(jiǎn)單來(lái)說(shuō)，大部分這些模型無(wú)法適應(yīng)單個(gè)GPU。

通常，GPU的內(nèi)存介于16到40GB之間，如果考慮一個(gè)30B模型，在半精度下需要60GB的內(nèi)存，或者70B Lama模型在半精度下至少需要140GB的內(nèi)存。這意味著至少需要8個(gè)GPU。因此，我們需要一種解決方案將這些模型分區(qū)到多個(gè)設(shè)備上。我們來(lái)看看在這個(gè)領(lǐng)域有哪些模型并行化的解決方案和方法。

目前有兩種主要方法。一種是張量并行，你基本上在op內(nèi)部（如矩陣乘法）上切割你的模型，從而并行化計(jì)算。這會(huì)引入一個(gè)通信，就像全歸約一樣，如果你有足夠的工作負(fù)載，使用流水線并行計(jì)算會(huì)更快，但需要更高速的網(wǎng)絡(luò)。原因是要保持GPU忙碌，因?yàn)樗窃诓僮鞑⑿谢小Ｋ运m用于單節(jié)點(diǎn)計(jì)算。

另一種主要方法是流水線并行，基本上將模型水平切分成多個(gè)階段，每個(gè)階段包含一組層。所以你將第一個(gè)階段的輸出作為第二個(gè)階段的輸入。在這里，我們需要將這種技術(shù)與micro批處理結(jié)合使用，否則無(wú)法有效利用GPU。由于通信的性質(zhì)，流水線并行更適用于多節(jié)點(diǎn)的擴(kuò)展。

·ModifyModelCode
Megatron,TransformerNeuronx,Fairscale,etc.
DefineParallellayers-buildthemodelwithparallellayers
boundtospecifictrainers
·PyTorchAPIs
WraptheoriginalmodelwithAPIs
Automaticallypartitionthemodel
Traineragnostic

那么，我們今天在OSS空間實(shí)際上如何實(shí)現(xiàn)這些模型呢？這里有兩種主要方法。一種是修改你的模型代碼。基本上，定義你的并行層，并在這些并行層之上構(gòu)建你的模型。這非常像transformers的風(fēng)格。我們?cè)诓煌膔epo中也看到了這一點(diǎn)，比如megatron,Transformer，Neuronics來(lái)自AWS。而且大多數(shù)這種并行性都限定在特定的訓(xùn)練上。而另一方面，我們有PyTorch的API，它采用了一種不同的策略。它大多數(shù)情況下不需要改變或只需進(jìn)行最小的改動(dòng)來(lái)適應(yīng)您的模型，它只需要檢查您的模型，并自動(dòng)對(duì)其進(jìn)行分割。它也可以在訓(xùn)練過(guò)程中保持中立。

這里有一個(gè)修改模型代碼的例子，這是在Fairscale上構(gòu)建的LLAMA2模型的例子。正如您在左側(cè)可以看到的那樣，我們實(shí)際上正在定義那些并行層，在注意層中，您可以看到我們正在在那些并行層之上構(gòu)建模型。如我所說(shuō)，這是Fairscale，但是Megatron，Transformer和AX都是相同的性質(zhì)。另一方面，這是PyTorch的API，用于流水線并行。我們有一個(gè)名為PP的軟件包，它處于測(cè)試階段。它的作用是使用模型并行地進(jìn)行計(jì)算

為了理解你的模型并將其分成多個(gè)階段，使用tracing方法是非常重要的。這個(gè)框架提供了一個(gè)非常簡(jiǎn)單的API，只需要將你的模型輸入其中，就能得到一個(gè)已經(jīng)在不同GPU上進(jìn)行了分布的多個(gè)階段。它還支持延遲初始化，我們稍后會(huì)詳細(xì)討論。此外，我們還有適用于PyTorch的tensor并行API，與張量一起使用。

如果你看一下這些代碼，基本上，你可以將你的分片策略傳遞給并行模塊，它會(huì)簡(jiǎn)單地將你的模塊并行化到不同的設(shè)備上。這又是相同的策略，你不需要改變你的模型代碼。它們都是訓(xùn)練無(wú)關(guān)的，所以你可以在從不同庫(kù)中導(dǎo)入任意的檢查點(diǎn)時(shí)進(jìn)行推理。接下來(lái)，我來(lái)強(qiáng)調(diào)一下我們?cè)诜植际酵评碇兴媾R的一些挑戰(zhàn)。首先，大多數(shù)的開(kāi)源解決方案都與特定的訓(xùn)練器綁定。

1. Most of the OS solutions are bound to specific trainers or require model changes：
-ExamplesincludeDeepSpeed,Accelerate,ParallelFormer,TGI,vLLM,etc.

2.Automaticpartitioningofarbitrarycheckpoints:
-APIstoautomaticallypartitionyourmodelcheckpoints(trainer-agnostic)
-Enablesautomaticpartitioningofarbitrarycheckpoints.

3.Deferredinitialization(loadingpretrainedweights):
-AvoidloadingthewholemodelonCPUordevice.
-Supportsdeferredinitializationtoloadpretrainedweightsondemand.

4.Checkpointconversion:
-Convertcheckpointstrained/savedwithdifferentecosystemlibrariestoPyTorchDistributedcheckpoints.

正如我提到的，他們需要模型的改變，比如deepspeed、VLLM等。所以這需要一種解決方案，能自動(dòng)對(duì)模型和任意檢查點(diǎn)進(jìn)行分區(qū)。所以無(wú)論你用哪個(gè)訓(xùn)練器訓(xùn)練過(guò)你的模型，它實(shí)際上應(yīng)該能夠?qū)⒛愕哪Ｐ瓦M(jìn)行分區(qū)。這里還有另外兩個(gè)挑戰(zhàn)，就是延遲初始化，正如Mark所談的。它可以幫助你更快地加載模型，并在某些情況下避免在CPU和GPU上的開(kāi)銷(xiāo)。而且，如果你必須將模型放在操作系統(tǒng)上，也可以使用這種方式。然后我們有一個(gè)檢查點(diǎn)轉(zhuǎn)換，我這里稍微談一下。這是今天的初始化方式。

所以你可以使用元設(shè)備來(lái)初始化你的模型。然后你請(qǐng)求模型并行API并行化你的模型。你必須手動(dòng)地實(shí)現(xiàn)你的參數(shù)。但是在這里還有一個(gè)額外的步驟，你必須將你的模型檢查點(diǎn)轉(zhuǎn)換為PyTorch分布式可以理解的張量形式。所以，在它們之間有一個(gè)檢查點(diǎn)轉(zhuǎn)換的過(guò)程。你可以使用PyTorch分布式檢查點(diǎn)API來(lái)加載模型。這樣，你實(shí)際上可以使用延遲初始化。這里有一個(gè)額外的步驟，我們正在研究如何在這里去除檢查點(diǎn)轉(zhuǎn)換。

好的，談?wù)撘幌路植际酵评砗筒煌哪Ｐ筒⑿谢，F(xiàn)在讓我們轉(zhuǎn)向Torchserve，看看我們?cè)赥orchserve上支持什么。今天在Torchserve上，我們已經(jīng)集成了分布式推理解決方案，我們與DeepSpeed、Hugging Face Accelerate、Neuron SDK與AWS自定義芯片都有集成。Torchserve原生API還具備PP和TP功能。我們還有微批處理、連續(xù)批處理和流式響應(yīng)的API，這是我們團(tuán)隊(duì)和AWS的Matias和Lee共同開(kāi)發(fā)的。

現(xiàn)在讓我們來(lái)看一下PP。在這里，我們有高度流水線并行處理，我們已經(jīng)初始化了這項(xiàng)工作，并啟用了路徑。我們的主要目標(biāo)是專(zhuān)注于使用的便利性和功能。您可以看到，我們?cè)谶@里提供了一個(gè)一行代碼的API，您可以輕松地將您的模型傳遞進(jìn)去，特別是所有的Hugging Face模型，您實(shí)際上可以獲得階段并簡(jiǎn)單地初始化您的模型。

然后，我們最近還為L(zhǎng)ama啟用了d-tensors來(lái)實(shí)現(xiàn)張量并行。這是初始的一步。我們已經(jīng)啟用了這條路徑，但我們也正在努力優(yōu)化推理路徑。所以請(qǐng)繼續(xù)關(guān)注這里，很快就會(huì)有新的更新。

關(guān)于微批處理，如果您想要使用管道并行處理，微批處理非常重要。我們有一些非常好的微批處理示例，可以幫助您輕松入門(mén)。它既有助于更好地利用GPU，也可以在某些情況下并行化預(yù)處理，比如處理一些您正在處理的視覺(jué)模型。所以可以有一個(gè)繁重的預(yù)處理任務(wù)，我們可以在這里實(shí)現(xiàn)并行化，使用多線程。我們還有連續(xù)批處理和其他供LLM服務(wù)配置使用的成分。這里的想法是當(dāng)一個(gè)請(qǐng)求完成時(shí)，將隊(duì)列中的請(qǐng)求連續(xù)添加到當(dāng)前批次中作為一個(gè)請(qǐng)求。所以你不需要等待整個(gè)批次完成再發(fā)送下一個(gè)請(qǐng)求。

正如馬克所說(shuō)的，基本上就是動(dòng)態(tài)批處理。因此，這將有助于提高吞吐量和用戶(hù)體驗(yàn)。我們來(lái)看下一個(gè)功能，即流式響應(yīng)API。

再次強(qiáng)調(diào)，當(dāng)您向這些LLMs發(fā)送請(qǐng)求時(shí)，它們可能需要很長(zhǎng)時(shí)間進(jìn)行推理和生成令牌。因此，流式API將幫助您獲取每個(gè)令牌的生成，而無(wú)需等待整個(gè)序列的生成。您將逐個(gè)生成的令牌返回到客戶(hù)端。因此，這是一個(gè)很好的功能可以實(shí)現(xiàn)和集成到友好的環(huán)境中。在這里，您可以看到我們定義了兩個(gè)API。一個(gè)是發(fā)送中間預(yù)測(cè)響應(yīng)的API，您可以使用該API。我們使用了HuggingFace文本迭代器來(lái)進(jìn)行流式批處理。通過(guò)這兩個(gè)的組合，我們?cè)谶@里實(shí)際上有LLAMA2的示例。

再說(shuō)一次，正如我所談到的，我們與所有這些功能進(jìn)行了集成，包括所有這些不同的庫(kù)，如HuggingFace、PP、DeepSpeed、DeepSpeedM2、Inferentia2。在這里，我們實(shí)際上已經(jīng)發(fā)布了一個(gè)很新的推理示例，并且我們?cè)诖税l(fā)布了一個(gè)使用案例。

你可以看到我們?cè)谶@里使用了張量并行和并行預(yù)填充。它實(shí)際上幫助大大加快了推理速度。與等效的GPU相比，它使得Inferentia 2的成本大幅降低了。我們獲得了3倍更便宜的性能點(diǎn)。因此，我強(qiáng)烈建議你也看一看這個(gè)示例。

Mark談了很多不同的優(yōu)化，補(bǔ)充優(yōu)化方法使模型更快。他談到了內(nèi)存限制、CPU限制，對(duì)于特定的LLMs來(lái)說(shuō)，還有兩個(gè)重要的事情，即KV緩存。這非常重要。它可能會(huì)占用非常多的內(nèi)存，并且會(huì)受到內(nèi)存限制的影響。因此，類(lèi)似于pageattention的想法在這里可能有幫助。另一個(gè)要考慮的因素是量化。

到此結(jié)束，謝謝。

審核編輯：黃飛