MaskLLM
NVIDIA Research · 程式設計 · 其他
MaskLLM 是 NVIDIA Research 推出的開源剪枝方法，靠可學習的 N:M 稀疏遮罩縮小大型語言模型，降低推論成本，同時把準確度守在原始模型附近。
關於 MaskLLM
What Is MaskLLM
MaskLLM 是一套用來壓縮大型語言模型的研究方法。它屬於剪枝工具的一員，負責移除模型用不到的參數，做法則採用半結構化稀疏，也就是所謂的 N:M 稀疏。道理很單純：在每 M 個連續權重組成的區塊裡，只留下 N 個。2:4 的樣式會在四個權重中保留兩個，等於把硬體需要搬動的資料量砍半。
MaskLLM 的真正差別，在於它怎麼決定留哪些權重。舊方法靠一條固定規則，例如丟掉最小的權重，或用統計估算挑選。MaskLLM 把遮罩的選擇當成模型要學的東西，將 N:M 樣式建構成機率分布，再透過一種叫 Gumbel Softmax 的取樣技巧訓練，讓模型能在大型資料集上找到更好的遮罩。學遮罩的過程中權重維持凍結，這也是為什麼對特定任務來說，2:4 稀疏能做到「無損」。
實際吸引力藏在數字裡。根據 NVIDIA 的說法，稠密的 LLaMA-2 7B 模型在 Wikitext 上的困惑度是 5.12，而先前的剪枝方法會把它推到 10 以上。MaskLLM 只靠學習遮罩、權重凍結不動，就把數字壓到 6.72 到 6.78。專案在 NVIDIA 授權下開源，研究人員和工程師都能自己跑。不過前提是要有真正的 GPU 硬體，還得懂機器學習。
Getting Started
- 安裝 Python 套件，再執行腳本，從 Hugging Face 下載已經算好的遮罩
- 在 PyTorch 裡載入目標模型，例如 LLaMA-2 7B 或 LLaMA-3 8B
- 套用下載的遮罩，並用專案附的困惑度腳本測準確度
- 若要自行訓練，架好 Megatron-LM 環境、備妥資料，再為自己的領域訓練一個新遮罩
產品資訊
快速了解 MaskLLM 的定價、支援平台與效能。
適合對象
這項工具最適合的使用者、任務與情境。
使用者
- 機器學習研究人員：研究模型壓縮的工程師，想要一套能執行、能擴充，還能跟 SparseGPT、Wanda 比較的可學習稀疏方法
- 部署模型的 ML 工程師：需要把大模型縮小到 GPU 預算吃得下、又不想重新訓練整個網路的團隊
- 研究生與學術工作者：做高效推論研究、需要開放原始碼與可重現數據的人
- 框架與工具開發者：想把 N
任務
- 大型語言模型剪枝：把模型剪到 2
- 一次性壓縮：不用長時間訓練，直接把算好的遮罩套到 LLaMA-2、LLaMA-3 或 LLaMA-3.1
- 領域專屬稀疏：為特定任務學一個專用遮罩，讓剪枝後的模型在該領域保住準確度
- 稀疏轉移：把學好的遮罩分布搬到不同領域或任務重複使用，不必從零開始
- 剪枝方法評比：在同一套設定下，把 MaskLLM 跟 SparseGPT、Wanda 和幅度剪枝放在一起比較
情境
- 團隊只有少數 GPU 卻要跑 7B 模型：用 2
- 研究需要重現論文數字：跑官方釋出的遮罩，對照論文檢查困惑度，例如 LLaMA-2 7B 的 6.78
- 模型整體表現不錯，但在利基領域失準：學一個量身打造的遮罩，讓剪枝版本在那些資料上依然準確
- 為一整個模型家族打造壓縮流程：先從既有遮罩出發，再逐步微調，最後裁切 checkpoint 準備上線
- 檢視授權與隱私限制：使用 NVIDIA 以自家授權釋出的原始碼與權重
主要功能
可學習剪枝的 N:M 遮罩
MaskLLM 不靠人工逐一評分權重，而是讓遮罩自己學出來。它把合法的 N:M 樣式集合當成一個分布，在訓練時從中取樣，模型便會收斂到最適合某份資料集的遮罩。這個可學習的步驟，正是它跟一次性方法的根本差別，後者選一次遮罩就停手。
Gumbel Softmax 取樣
這套方法借用了離散取樣的技巧，讓遮罩的選擇可以端到端訓練。Gumbel Softmax 讓梯度能穿過原本非黑即白的硬選擇，也因為這樣，N:M 遮罩上的分布才有辦法用真實資料最佳化，而不是靠估算。
現成的 Hugging Face 模型遮罩
專案附上現成遮罩，不用訓練就能試。官方釋出的遮罩涵蓋 LLaMA-2 7B、LLaMA-3 8B 與 LLaMA-3.1 8B，各以 C4 資料集的子集訓練而成，並有腳本自動從 Hugging Face 下載、套用到官方模型上評估。
半結構化 2:4 稀疏
MaskLLM 瞄準的是 N:M 稀疏，多半是每四個連續權重保留兩個非零值的 2:4 樣式。這個格式之所以重要，是因為它能對上硬體加速，所以縮小的模型不只佔更少空間，在對的 GPU 上還能跑得更快。
稀疏樣式轉移
同一套學遮罩的機率模型，可以把遮罩從一個領域帶到另一個。這種遮罩轉移在分布訓練好之後，要用到新任務或新資料集時，所需力氣比從零開始少，作者稱之為可轉移性。當一個模型得在好幾種情境都保持精準時，這一點特別有用。
Megatron-LM 訓練流程
要跑完整訓練時，MaskLLM 備有 Megatron-LM 工作流程，把 LLaMA 模型分散到多張 GPU 上。專案提供腳本，用來轉換 checkpoint、將資料分詞，用 SparseGPT、Wanda 或幅度剪枝產生先驗遮罩，接著訓練遮罩，最後在發布前裁切 checkpoint。
優缺點
優點
- 從資料中學遮罩，而不是死守固定的重要性規則，同樣稀疏度下準確度更好
- 以開放授權釋出原始碼、權重與現成遮罩，成果可以重現
- 支援遷移學習，一份遮罩分布可用於多個任務與領域
- 透過 Hugging Face 與 Megatron-LM，可搭配 LLaMA-2、LLaMA-3、LLaMA-3.1 等常見模型
- 權重維持凍結，學到的 2:4 稀疏對調校過的任務而言沒有損失
缺點
- 訓練新遮罩需要多 GPU 機器，專案指出完整訓練每張 GPU 約需 40GB 記憶體
- 它是研究產物，不是產品，因此沒有客服、沒有代管服務、也沒有 API
- 現成遮罩只涵蓋少數 LLaMA 模型，其他架構得自己訓練
- 效益出現在支援 2:4 加速的硬體上，老舊或非 NVIDIA 的 GPU 能得到的幫助有限
常見問題
MaskLLM 是一套可學習的剪枝方法，為大型語言模型加上 N:M 稀疏。它不逐一評分權重，而是把遮罩建構成機率分布，從中學出該保留哪些權重。NVIDIA Research 在 NeurIPS 2024 發表，MaskLLM 原始碼在 GitHub 上開放。
相關內容
探索與 MaskLLM 相關的工具、技能與文章。
MaskLLM 替代方案
Base Code
Base44 (Wix) · 程式設計 · 生產力Base Code 是 Base44 推出的共享雲端開發環境。團隊在對話中描述對既有 GitHub 儲存庫的修改，看著預覽確認結果，再以 Pull Request 交回工程團隊。
Google EmbeddingGemma 2
Google DeepMind · 寫作 · 圖片 · 程式設計Google EmbeddingGemma 2 是 7.4 億參數的開源模型，能把文字、程式碼、圖片、音訊與影片畫面轉成同一個 768 維空間裡的向量，在手機與筆電上就能執行。
GitGlow
GitGlow · 程式設計 · 生產力GitGlow 是桌面版 Git 用戶端，把你自己的程式碼與 AI 代理程式的變更集中在同一份審查裡，讓你在交付前看清楚機器上究竟要送出什麼。