大模型雖然好，但我的筆記本和手機(jī)都跑不動呀。就算勉強(qiáng)能跑起來，也是奇慢無比。而與此同時(shí)，對適合移動和邊緣設(shè)備的小模型的需求卻在不斷增長，因?yàn)檫@些模型似乎才能真正滿足人們的日常需求。正因?yàn)榇?，有不少研究者和?yīng)用開發(fā)者都認(rèn)為小模型才是 AI 的未來。

事實(shí)上，Meta 和 Mistral 等都已經(jīng)發(fā)布了自己的 SLM，比如 Llama 3.2 的 1B 和 3B 版本以及 Ministral 3B。另外還有一些社區(qū)開發(fā)的 SLM，比如 BabyLlama 系列（不到 1B 參數(shù)）、 TinyLLaMA（1.1B 參數(shù)）。

為了打造出真正好用的小型語言模型（SLM），AI 研究社區(qū)想出了各種各樣的方法，像是對大模型進(jìn)行蒸餾或量化或者就直接去訓(xùn)練性能優(yōu)異的小模型。

實(shí)際上， SLM 正在逐漸成為一個(gè)研究熱門方向，簡單檢索 arXiv 上的關(guān)鍵詞也能大致看見這一趨勢：9 和 10 月份，SLM 相關(guān)研究論文的數(shù)量有了明顯增長。

今天，我們就來看看蘋果的一篇相關(guān)論文，其探討了訓(xùn)練小型語言模型的計(jì)算瓶頸。

• 論文標(biāo)題：Computational Bottlenecks of Training Small-scale Large Language Models
• 論文地址：https://arxiv.org/pdf/2410.19456

首先，多小的模型才能算是小型語言模型，或按蘋果的說法 —— 小規(guī)模大型語言模型？

這個(gè)蘋果團(tuán)隊(duì)給出的指標(biāo)是「參數(shù)量 ≤ 2B」。當(dāng)然，這并非人們公認(rèn)的標(biāo)準(zhǔn)，也有人認(rèn)為 Ministral 3B 和 Llama 3.2 3B 等 3B 參數(shù)量的模型也算是 SLM?？傊?，大與小是一個(gè)會隨著計(jì)算基礎(chǔ)設(shè)施的演進(jìn)而動態(tài)變化的標(biāo)準(zhǔn)，昨天的大模型可能就會成為明天的小模型。

盡管 SLM 規(guī)模很小，但其表現(xiàn)并不一定很差，并且已經(jīng)展現(xiàn)出了自己的巨大潛力。很多借助剪枝、蒸餾和量化等技術(shù)得到的 SLM 的性能并不比大得多的模型差，甚至有時(shí)候還能更勝一籌。舉個(gè)例子，Gemma-2B 的性能就優(yōu)于大得多的 OPT-175B，這就挑戰(zhàn)了大多數(shù)人的一個(gè)固有觀念：模型大小是有效性的主導(dǎo)決定因素。

另外，也有采用互相驗(yàn)證等其它新穎方法提升 SLM 能力的研究思路，比如機(jī)器之心曾報(bào)道過的《兩個(gè)小模型互相驗(yàn)證，直接比肩大模型？微軟的 rStar 甚至沒用 CoT 和微調(diào)》。事實(shí)上，隨著 OpenAI ο1 系列模型的發(fā)布，通過優(yōu)化推理時(shí)間計(jì)算也成了提升 SLM 性能的重要途徑。

性能足夠好的 SLM 具有很大的好處，最基本的就是速度快、效率高、性價(jià)比高。因此，SLM 對計(jì)算資源有限的組織（如小型企業(yè)和學(xué)術(shù)機(jī)構(gòu)）非常有吸引力。

蘋果的這項(xiàng)研究關(guān)注的是 SLM 的訓(xùn)練動態(tài)。事實(shí)上，在訓(xùn)練方面，LLM 和 SLM 的差距很大。LLM 的計(jì)算需求和基礎(chǔ)設(shè)施需求并不一定適用于 SLM。考慮到云平臺可用的硬件配置多種多樣（包括 GPU 類型、批量大小和通信協(xié)議），有必要對這些影響 SLM 訓(xùn)練效率的因素進(jìn)行系統(tǒng)性的分析，尤其是要考慮一些符合實(shí)際的指標(biāo)，比如每美元的損失和每秒 token 數(shù)。

該團(tuán)隊(duì)的研究結(jié)果表明，對于更小型的模型，可以使用 A100-40GB GPU 和分布式數(shù)據(jù)并行（DDP）等更低成本選擇，同時(shí)不會對性能產(chǎn)生負(fù)面影響。對于更大型的模型，就必需更高級的配置了（例如 A100-80GB 和 H100-80GB GPU 搭配 Flash Attention（FA）和完全分片式數(shù)據(jù)并行（FSDP）），這樣才能處理更大的數(shù)據(jù)批以及防止內(nèi)存相關(guān)的問題。

SLM 領(lǐng)域的最近研究進(jìn)展表明，擴(kuò)展 AI 系統(tǒng)不僅是要追求先進(jìn)的性能，也要考慮實(shí)際應(yīng)用。目前這股研發(fā) SLM 的趨勢表明，重新評估硬件和計(jì)算策略是非常重要的。

蘋果這項(xiàng)研究為此做出了貢獻(xiàn)，他們系統(tǒng)性地研究了在不同的云基礎(chǔ)設(shè)施和設(shè)置上，訓(xùn)練最多 2B 參數(shù)大小的 SLM 的計(jì)算瓶頸和成本效率。他們發(fā)現(xiàn)：

1. 相比于 LLM，F(xiàn)lashAttention 對 SLM 來說更重要；

2. H100-80GB 和 A100-80GB 等昂貴硬件對 SLM 訓(xùn)練來說不一定具有成本效益；

3. DDP 是 SLM 的最佳分布式訓(xùn)練方案；

4. 對 SLM 訓(xùn)練來說，最大化 GPU 內(nèi)存利用率并不是成本最優(yōu)的。

模型和參數(shù)

該團(tuán)隊(duì)研究的是 LLaMa 架構(gòu)，畢竟不管是 LLM 還是 SLM，這都是當(dāng)今最流行的架構(gòu)。

LLaMa-2 和 3 最小的版本分別是 7B 和 8B，但這對大多數(shù)移動硬件來說還是太大了。為此，該團(tuán)隊(duì)進(jìn)行了一番操作：為了定義他們自己的模型，他們通過在 Llama 模型上擬合一條曲線而提取了模型的解碼器模塊和參數(shù)數(shù)量。見下圖 5。

他們評估了四種不同的模型大小：100M、500M、1B 和 2B。

值得注意的是，他們最大化了圖中 x 軸或圖例中未顯示的所有配置參數(shù)。也就是說，他們對下面列出的所有配置參數(shù)組合進(jìn)行大型網(wǎng)格搜索，并且每個(gè)圖中的每個(gè)點(diǎn)都是給定圖中指定的所有參數(shù)的最佳配置。

這樣，他們找到了最佳的 Token/Dollar 比值，并假設(shè)可以通過調(diào)整優(yōu)化超參數(shù)（例如學(xué)習(xí)率）來實(shí)現(xiàn)與硬件最佳配置的最佳收斂。

接下來，給出這些配置參數(shù)的定義

• GPU 類型：他們評估了三種英偉達(dá) GPU：A100-40GB、A100-80GB 和 H100-80GB。所有 GPU 使用的數(shù)據(jù)類型都是 BFloat16。
• GPU 數(shù)量和通信：每種 GPU 類型都有三種主要訓(xùn)練配置，包括：單節(jié)點(diǎn)單 GPU（1 臺 GPU）、單節(jié)點(diǎn)多 GPU（2、4 和 8 臺 GPU）和多節(jié)點(diǎn)多 GPU（16、32 和 64 臺 GPU）。當(dāng) GPU 超過 1 臺時(shí)，還會評估分布式數(shù)據(jù)并行（DDP）和完全分片式數(shù)據(jù)并行（FSDP）這兩種通信方法。對于分片，他們也研究了兩種策略：1) 完全分片，即對所有梯度、優(yōu)化器狀態(tài)和權(quán)重進(jìn)行分片；2) grad_op 分片，即僅對梯度和優(yōu)化器狀態(tài)進(jìn)行分片（但權(quán)重不分片）。他們使用了 RDMA/EFA。
• 樣本數(shù)量：他們還評估了訓(xùn)練期間適合單個(gè) GPU 的各種樣本數(shù)量。他們將序列長度固定為 1028，并迭代適合單臺設(shè)備的批量大小。由于即使在最?。?00M）的模型中也無法將 128 個(gè)樣本放入單個(gè) GPU 內(nèi)存中，因此他們研究了每臺設(shè)備的批量大小為 4、8、16、32 和 64 的情況。這里不使用梯度累積
• Flash Attention：他們研究了在注意力模塊中使用 Flash Attention 的影響。

實(shí)驗(yàn)結(jié)果

該團(tuán)隊(duì)展示了 SLM 在 A100-40GB、A100-80GB 和 H100-80GB 上運(yùn)行的結(jié)果。他們在 HuggingFace 中部署了模型，沒有附加任何額外的框架，batch size=1024，并使用 PyTorch 運(yùn)行了如下表數(shù)值設(shè)置的實(shí)驗(yàn)。

通過多次實(shí)驗(yàn)，研究團(tuán)隊(duì)得出了以下觀察：

Q1：在 SLM 訓(xùn)練中使用 FlashAttention 有多重要？

圖 1 比較了不同 batch size 下 FlashAttention2 與普通注意力的區(qū)別。首先，使用 FlashAttention 顯著提高了 SLM 的 Token/Dollar 效率，用同樣的錢（Dollar）能處理更多數(shù)據(jù)（Token）。

對于較小的模型，F(xiàn)A 對 Token/Dollar 的提升更為顯著。這是因?yàn)樽⒁饬C(jī)制的成本會隨上下文長度變長以平方的速度增長。當(dāng)模型的隱藏層維度減少時(shí)，這個(gè)因素變得尤為重要。這樣的 SLM，其性能受限于數(shù)據(jù)處理能力，其中數(shù)據(jù)在 CPU/GPU、GPU/GPU 之間傳輸是主要瓶頸。最后，可以看到對于較大的模型（1B 和 2B），F(xiàn)A 能夠訓(xùn)練更大的 batch size（1024），而普通注意力會導(dǎo)致內(nèi)存溢出（OOM）。

Q2：GPU 數(shù)量一定，哪類 GPU 最適合訓(xùn)練 SLM？

圖 2 展示了使用 A100-40GB 和 A100-80GB GPU 訓(xùn)練模型的效果。盡管不同模型間沒有統(tǒng)一的趨勢，但當(dāng)使用用大量 GPU 訓(xùn)練 1B、2B 參數(shù)規(guī)模的模型時(shí)，A100-80GB GPU 表現(xiàn)更佳。這種 GPU 適合處理更大的 batch size。對于更小的模型，則可以選擇成本更低的 40GB GPU。

Q3：在不同節(jié)點(diǎn)數(shù)量下，哪種通信方案最適合訓(xùn)練 SLM？

該團(tuán)隊(duì)探討了不同并行策略對 SLM 訓(xùn)練的影響，分布式數(shù)據(jù)并行（DDP）、完全分片數(shù)據(jù)并行（FSDP-Full）、還是 FSDPGrad+Optimizer，那種方法更優(yōu)秀？圖 3 展示了在 A100-80GB GPU 上，這些并行策略訓(xùn)練模型的效果。

結(jié)果顯示，對于小型模型，對通信需求最小的 DDP 更優(yōu)。但對于 2B 參數(shù)的模型，F(xiàn)SDP 由于能處理更大的 batch size，表現(xiàn)超越了 DDP。此外，F(xiàn)SDP-Grad+Optimizer 因其較低的通信開銷，表現(xiàn)優(yōu)于 FSDP-Full。簡而言之，選擇合適的并行策略可以優(yōu)化 SLM 的訓(xùn)練效率。

Q4：對于不同的全局 batch size，訓(xùn)練 SLM 的最佳通信方案是什么？

圖 4 顯示了使用 DDP、FSDP-Full 和 FSDP-Grad+Optimizer 在不同 batch size 下訓(xùn)練 SLM 的結(jié)果。

可以看到，batch size 較小時(shí)，通信方案的差異不大。然而，與上個(gè)問題類似，對于 2B 模型和 batch size，F(xiàn)SDP 方案的性能優(yōu)于 DDP，并且能夠處理比 DDP 更大的批量大小，而 DDP 在這種情況下會導(dǎo)致內(nèi)存溢出。

更多研究細(xì)節(jié)，請?jiān)L問原論文。