GaLore：推進消費級硬體上的大模型訓練
GaLore 在大型語言模型 (LLM) 訓練中的整合標誌著深度學習領域的一項重大進步，特別是在記憶體效率和人工智慧研究的民主化方面。透過在消費級硬體上訓練數十億引數的模型，減少最佳化器狀態的記憶體佔用，並利用先進的投影矩陣技術，GaLore 為計算資源有限的研究人員和從業者開闢了新的視野。
利用消費級硬體擴充套件LLM
GaLore 在消費級 GPU（如 NVIDIA RTX 4090）上訓練多達 70 億引數模型（如基於 Llama 架構的模型）的能力是開創性的。這透過顯著降低訓練過程中最佳化器狀態和梯度所需的記憶體要求來實現。該方法利用深度神經網路中梯度固有的低秩結構，應用投影來降低需要儲存和操作的資料的維度。
最佳化器狀態的記憶體效率
最佳化器狀態，特別是在 Adam 等自適應最佳化演算法中，在模型訓練過程中佔據了記憶體佔用的大部分。GaLore 透過在梯度被最佳化器處理之前將其投影到較低維的子空間來解決這個問題。這不僅減少了儲存這些狀態所需的記憶體，還保持了最佳化過程的有效性。
記憶體節省是可觀的，作者報告“**在訓練過程中，最佳化器狀態的記憶體減少了 82.5% 以上**”，使得在相同記憶體限制下訓練更大的模型或使用更大的批次大小成為可能。當與 8 位精度最佳化器結合使用時，這些節省甚至會更加顯著。
子空間切換和高階投影技術
GaLore 有效性的一個關鍵組成部分是其動態子空間切換機制，該機制允許模型在整個訓練過程中在不同的低秩子空間中導航。這確保了模型不會侷限於引數空間的有限部分，從而保留了全引數學習的能力。何時以及如何切換子空間的決定至關重要，這些切換的頻率是保持一致的最佳化軌跡和適應梯度低秩結構不斷演變的平衡。
動態調整這些投影以響應梯度結構變化的能力是 GaLore 武器庫中的一個強大工具，它允許更細緻地控制訓練大型模型固有的記憶體最佳化權衡。
GaLore 與 8 位最佳化器結合
GaLore 與 8 位精度最佳化器的結合代表了一種協同效應，它最大限度地提高了記憶體效率，同時保持了訓練過程的完整性和效能。8 位最佳化器透過量化最佳化器狀態來減少記憶體佔用。當與 GaLore 的投影機制結合使用時，結果是一個高度記憶體高效的訓練方案，而不會損害模型精度或收斂速度。
這種組合在記憶體是關鍵瓶頸的場景中特別有效，例如在消費級硬體上訓練大型模型或在記憶體受限的環境中部署模型。它可以在相同的硬體限制下使用更復雜的模型和更大的資料集，從而突破有限資源所能實現的界限。
實現細節
將 8 位最佳化器與 GaLore 整合以訓練大型語言模型 (LLM) 涉及將梯度、權重和最佳化器狀態量化為 8 位表示。此量化過程顯著減少了記憶體佔用，從而可以在相同記憶體限制下訓練更大的模型或使用更大的批次大小。這種整合的演算法細節涉及幾個關鍵步驟，其中一些將受益於原生 CUDA 實現以提高效率。GaLore 為將這些技術更緊密地與量化和矩陣的專門引數化相結合開闢了新的可能性，這可以進一步減少記憶體使用。我們目前正在 bitsandbytes 庫中探索這個方向。
GaLore 8 位最佳化演算法概述
梯度投影：GaLore 使用投影矩陣將全精度梯度投影到低秩子空間中。此步驟減少了梯度的維度，然後將其量化為 8 位格式。
量化：投影梯度以及模型權重和最佳化器狀態（例如 Adam 中的移動平均值）從 32 位浮點數量化為 8 位整數表示。這涉及將浮點值縮放到 8 位範圍並將其四捨五入到最接近的整數。
最佳化器更新：8 位量化梯度用於更新模型權重。此步驟涉及將梯度反量化回浮點格式，應用最佳化器的更新規則（例如，Adam 的動量更新和引數調整），然後將更新的最佳化器狀態量化回 8 位以進行儲存。
反量化和權重更新：8 位量化權重經過反量化為浮點表示以進行處理，儘管由於值的有限範圍，它們保留了量化形式固有的 8 位精度。此步驟是必需的，因為 PyTorch 等框架中的標準操作不支援 8 位整數，並且此類整數權重無法容納梯度。雖然此方法本質上不提高準確性，但它有助於在當前深度學習庫的限制內實際應用和計算量化權重的梯度。請注意，在反量化和應用權重更新之後，GaLore 會再使用一次投影，將反量化的低秩更新投影回原始空間。
與 Hugging Face Transformers 一起使用
要將 GaLore 最佳化器與 Hugging Face transformers 庫一起使用，您需要先將其更新到支援 GaLore 最佳化器的版本，可以透過安裝最新更新，即 `pip install transformers>=4.39.0`，或從原始碼安裝 transformers。
然後使用 `pip install galore-torch` 安裝 galore-torch 庫。下面是 GaLore 與 transformers 配合使用的完整示例，用於在 imdb 資料集上預訓練 Mistral-7B。
import torch
import datasets
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
import trl
train_dataset = datasets.load_dataset('imdb', split='train')
args = TrainingArguments(
output_dir="./test-galore",
max_steps=100,
per_device_train_batch_size=2,
optim="galore_adamw",
optim_target_modules=["attn", "mlp"]
)
model_id = "mistralai/Mistral-7B-v0.1"
config = AutoConfig.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)
trainer = trl.SFTTrainer(
model=model,
args=args,
train_dataset=train_dataset,
dataset_text_field='text',
max_seq_length=512,
)
trainer.train()
`TrainingArguments`：只需傳遞一個有效的 `optim_target_modules`（它支援單個字串、正則表示式或字串或正則表示式列表），以及一個有效的 GaLore 最佳化器（例如 `galore_adamw`、`galore_adamw_8bit`、`galore_adafactor`）作為 `optim`，即可開始使用！
逐層更新
另一個需要提及的重要點是*逐層*最佳化器（即一次只更新一層權重）。通常，最佳化器在反向傳播後對所有層執行一次權重更新。這透過在記憶體中儲存整個權重梯度來完成。透過採用逐層權重更新，我們可以進一步減少訓練期間的記憶體佔用。在底層，這是透過 PyTorch 在使用者想要更新的層上使用後累積鉤子實現的。
要使用此功能，只需在最佳化器名稱後附加 `_layerwise`，例如 `galore_adamw_layerwise`。
結論
GaLore 憑藉其利用梯度低秩結構的創新方法，代表了 LLM 記憶體高效訓練方面向前邁出的重要一步。透過在消費級硬體上訓練數十億引數的模型，透過投影技術減少最佳化器狀態的記憶體佔用，並允許動態子空間切換，GaLore 民主化了大規模模型訓練的訪問。GaLore 與 8 位精度最佳化器的相容性進一步增強了其實用性，為在不需要專門計算資源的情況下訓練更大、更復雜的模型提供了途徑。這為人工智慧的研究和應用開闢了新的可能性，使其成為從業者和研究人員都感到興奮的時刻。
資源
請參閱原始論文。Twitter 參考：1 2 3。該論文還比較了 GaLore 和 ReLoRA，這可能引起一些讀者的興趣。對於在閱讀論文後仍有疑問或希望建設性地討論結果的讀者，請隨時加入作者的 Slack 社群。對於對類似釋出感興趣的讀者，請關注Jiawei Zhao和Titus von Koeller（獲取最新 `bitsandbytes` 版本資訊）以及Younes Belkada（獲取 Hugging Face 生態系統內外的量化相關最新資訊）。
