大模型已過時，小模型SLM才是未來？蘋果正在研究這個

大模型雖然好，但我的筆記本和手機都跑不動呀。 就算勉強能跑起來，也是奇慢無比。 而與此同時，對適合移動和邊緣裝置的小模型的需求卻在不斷增長，因為這些模型似乎才能真正滿足人們的日常需求。 正因為此，有不少研究者和應用開發者都認為小模型才是 AI 的未來。

事實上，Meta 和 Mistral 等都已經發布了自己的 SLM，比如 Llama 3.2 的 1B 和 3B 版本以及 Ministral 3B。另外還有一些社羣開發的 SLM，比如 BabyLlama 系列（不到 1B 引數）、 TinyLLaMA（1.1B 引數）。

為了打造出真正好用的小型語言模型（SLM），AI 研究社羣想出了各種各樣的方法，像是對大模型進行蒸餾或量化或者就直接去訓練效能優異的小模型。

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

今天，我們就來看看蘋果的一篇相關論文，其探討了訓練小型語言模型的計算瓶頸。

論文標題：Computational Bottlenecks of Training Small-scale Large Language Models

論文地址：https://arxiv.org/pdf/2410.19456

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

這個蘋果團隊給出的指標是「引數量 ≤ 2B」。當然，這並非人們公認的標準，也有人認為 Ministral 3B 和 Llama 3.2 3B 等 3B 引數量的模型也算是 SLM。總之，大與小是一個會隨著計算基礎設施的演進而動態變化的標準，昨天的大模型可能就會成為明天的小模型。

儘管 SLM 規模很小，但其表現並不一定很差，並且已經展現出了自己的巨大潛力。很多借助剪枝、蒸餾和量化等技術得到的 SLM 的效能並不比大得多的模型差，甚至有時候還能更勝一籌。舉個例子，Gemma-2B 的效能就優於大得多的 OPT-175B，這就挑戰了大多數人的一個固有觀念：模型大小是有效性的主導決定因素。

另外，也有采用互相驗證等其它新穎方法提升 SLM 能力的研究思路，比如機器之心曾報道過的《兩個小模型互相驗證，直接比肩大模型？微軟的 rStar 甚至沒用 CoT 和微調》。事實上，隨著 OpenAI ο1 系列模型的釋出，透過最佳化推理時間計算也成了提升 SLM 效能的重要途徑。

效能足夠好的 SLM 具有很大的好處，最基本的就是速度快、效率高、價效比高。因此，SLM 對計算資源有限的組織（如小型企業和學術機構）非常有吸引力。

蘋果的這項研究關注的是 SLM 的訓練動態。事實上，在訓練方面，LLM 和 SLM 的差距很大。LLM 的計算需求和基礎設施需求並不一定適用於 SLM。考慮到雲平臺可用的硬體配置多種多樣（包括 GPU 型別、批次大小和通訊協議），有必要對這些影響 SLM 訓練效率的因素進行系統性的分析，尤其是要考慮一些符合實際的指標，比如每美元的損失和每秒 token 數。

該團隊的研究結果表明，對於更小型的模型，可以使用 A100-40GB GPU 和分散式資料並行（DDP）等更低成本選擇，同時不會對效能產生負面影響。對於更大型的模型，就必需更高階的配置了（例如 A100-80GB 和 H100-80GB GPU 搭配 Flash Attention（FA）和完全分片式資料並行（FSDP）），這樣才能處理更大的資料批以及防止記憶體相關的問題。

SLM 領域的最近研究進展表明，擴充套件 AI 系統不僅是要追求先進的效能，也要考慮實際應用。目前這股研發 SLM 的趨勢表明，重新評估硬體和計算策略是非常重要的。

蘋果這項研究為此做出了貢獻，他們系統性地研究了在不同的雲基礎設施和設定上，訓練最多 2B 引數大小的 SLM 的計算瓶頸和成本效率。他們發現：

1. 相比於 LLM，FlashAttention 對 SLM 來說更重要；

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

3. DDP 是 SLM 的最佳分散式訓練方案；

4. 對 SLM 訓練來說，最大化 GPU 記憶體利用率並不是成本最優的。

模型和引數

該團隊研究的是 LLaMa 架構，畢竟不管是 LLM 還是 SLM，這都是當今最流行的架構。

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

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

值得注意的是，他們最大化了圖中 x 軸或圖例中未顯示的所有配置引數。也就是說，他們對下面列出的所有配置引數組合進行大型網格搜尋，並且每個圖中的每個點都是給定圖中指定的所有引數的最佳配置。

這樣，他們找到了最佳的 Token/Dollar 比值，並假設可以透過調整最佳化超引數（例如學習率）來實現與硬體最佳配置的最佳收斂。

接下來，給出這些配置引數的定義

GPU 型別：他們評估了三種英偉達 GPU：A100-40GB、A100-80GB 和 H100-80GB。所有 GPU 使用的資料型別都是 BFloat16。

GPU 數量和通訊：每種 GPU 型別都有三種主要訓練配置，包括：單節點單 GPU（1 臺 GPU）、單節點多 GPU（2、4 和 8 臺 GPU）和多節點多 GPU（16、32 和 64 臺 GPU）。當 GPU 超過 1 臺時，還會評估分散式資料並行（DDP）和完全分片式資料並行（FSDP）這兩種通訊方法。對於分片，他們也研究了兩種策略：1) 完全分片，即對所有梯度、最佳化器狀態和權重進行分片；2) grad_op 分片，即僅對梯度和最佳化器狀態進行分片（但權重不分片）。他們使用了 RDMA/EFA。

樣本數量：他們還評估了訓練期間適合單個 GPU 的各種樣本數量。他們將序列長度固定為 1028，並迭代適合單臺裝置的批次大小。由於即使在最小（100M）的模型中也無法將 128 個樣本放入單個 GPU 記憶體中，因此他們研究了每臺裝置的批次大小為 4、8、16、32 和 64 的情況。這裡不使用梯度累積

Flash Attention：他們研究了在注意力模組中使用 Flash Attention 的影響。

實驗結果

該團隊展示了 SLM 在 A100-40GB、A100-80GB 和 H100-80GB 上執行的結果。他們在 HuggingFace 中部署了模型，沒有附加任何額外的框架，batch size=1024，並使用 PyTorch 執行了如下表數值設定的實驗。

透過多次實驗，研究團隊得出了以下觀察：

Q1：在 SLM 訓練中使用 FlashAttention 有多重要？

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

對於較小的模型，FA 對 Token/Dollar 的提升更為顯著。這是因為注意力機制的成本會隨上下文長度變長以平方的速度增長。當模型的隱藏層維度減少時，這個因素變得尤為重要。這樣的 SLM，其效能受限於資料處理能力，其中資料在 CPU/GPU、GPU/GPU 之間傳輸是主要瓶頸。最後，可以看到對於較大的模型（1B 和 2B），FA 能夠訓練更大的 batch size（1024），而普通注意力會導致記憶體溢位（OOM）。

Q2：GPU 數量一定，哪類 GPU 最適合訓練 SLM？

圖 2 展示了使用 A100-40GB 和 A100-80GB GPU 訓練模型的效果。儘管不同模型間沒有統一的趨勢，但當使用用大量 GPU 訓練 1B、2B 引數規模的模型時，A100-80GB GPU 表現更佳。這種 GPU 適合處理更大的 batch size。對於更小的模型，則可以選擇成本更低的 40GB GPU。

Q3：在不同節點數量下，哪種通訊方案最適合訓練 SLM？

該團隊探討了不同並行策略對 SLM 訓練的影響，分散式資料並行（DDP）、完全分片資料並行（FSDP-Full）、還是 FSDPGrad+Optimizer，那種方法更優秀？圖 3 展示了在 A100-80GB GPU 上，這些並行策略訓練模型的效果。

結果顯示，對於小型模型，對通訊需求最小的 DDP 更優。但對於 2B 引數的模型，FSDP 由於能處理更大的 batch size，表現超越了 DDP。此外，FSDP-Grad+Optimizer 因其較低的通訊開銷，表現優於 FSDP-Full。簡而言之，選擇合適的並行策略可以最佳化 SLM 的訓練效率。

Q4：對於不同的全域性 batch size，訓練 SLM 的最佳通訊方案是什麼？

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

可以看到，batch size 較小時，通訊方案的差異不大。然而，與上個問題類似，對於 2B 模型和 batch size，FSDP 方案的效能優於 DDP，並且能夠處理比 DDP 更大的批次大小，而 DDP 在這種情況下會導致記憶體溢位。

更多研究細節，請訪問原論文。

參考連結：https://arxiv.org/pdf/2410.19456