大模型雖然好，但我的筆記本和手機都跑不動呀。就算勉強能跑起來，也是奇慢無比。而與此同時，對適合移動和邊緣設備的小模型的需求卻在不斷增長，因為這些模型似乎才能真正滿足人們的日常需求。正因為此，有不少研究者和應用開發者都認為小模型才是 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 在這種情況下會導致內存溢出。

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