8 位最佳化器

有狀態的最佳化器會隨時間維護梯度統計資訊，例如，歷史梯度值的指數平滑和（帶動量的 SGD）或平方和（Adam）。與普通的隨機梯度下降相比，這種狀態可以用來加速最佳化，但它會佔用本可以分配給模型引數的記憶體。因此，這限制了實踐中可訓練模型的最大尺寸。現在來看一下使用 8 位最佳化器可以訓練的最大模型。

bitsandbytes 最佳化器使用 8 位統計資訊，同時保持與使用 32 位最佳化器狀態相同的效能水平。

為了克服由此帶來的計算、量化和穩定性挑戰，8 位最佳化器包含三個部分

1. 分塊量化：將輸入張量分成更小的塊，並獨立量化，從而隔離異常值並將誤差更均勻地分佈到所有位上。每個塊在核心之間並行處理，從而實現更快的最佳化和高精度的量化。
2. 動態量化：高精度地量化小值和大值。
3. 穩定的嵌入層：在最佳化帶有詞嵌入的模型時提高穩定性。

有了這些元件，使用 8 位狀態執行最佳化器更新就變得很簡單。在執行更新之前，8 位最佳化器狀態會反量化為 32 位，然後狀態會再次量化回 8 位進行儲存。

8 位到 32 位的轉換在暫存器中逐元素進行，這意味著執行量化和反量化不需要緩慢地複製到 GPU 記憶體，也不需要額外的臨時記憶體。對於 GPU 而言，這使得 8 位最佳化器比常規的 32 位最佳化器快得多。

穩定的嵌入層

穩定的嵌入層提高了標準詞嵌入層在自然語言處理任務中的訓練穩定性。它解決了輸入分佈不均勻的挑戰，並減輕了極端的梯度變化。這意味著穩定的嵌入層可以支援更激進的量化策略而不會影響訓練穩定性，並且有助於實現穩定的訓練結果，這對於處理多樣化和複雜語言資料的模型尤其重要。

穩定的嵌入層有三個特點

• 初始化：利用 Xavier 均勻初始化來保持一致的方差，從而降低出現大梯度的可能性。
• 歸一化：在新增位置嵌入之前加入了層歸一化，有助於輸出的穩定性。
• 最佳化器狀態：該層專用的 32 位最佳化器狀態以增強穩定性，而模型的其餘部分可能使用標準的 16 位精度。

分頁最佳化器

分頁最佳化器是基於 CUDA 的統一記憶體功能構建的。統一記憶體提供了一個 GPU 和 CPU 可以輕鬆訪問的單一記憶體空間。雖然 PyTorch 不支援此功能，但它已被新增到 bitsandbytes 中。

分頁最佳化器的工作方式類似於常規的 CPU 分頁，這意味著它只有在 GPU 記憶體耗盡時才會啟用。當發生這種情況時，記憶體會逐頁從 GPU 轉移到 CPU。記憶體是對映的，這意味著頁面在 CPU 上是預先分配的，但不會自動更新。只有在訪問記憶體或啟動交換操作時，頁面才會被更新。

統一記憶體功能的效率低於常規的非同步記憶體傳輸，您通常無法獲得完整的 PCIe 記憶體頻寬利用率。如果您進行手動預取，傳輸速度可能會很高，但仍然只有完整 PCIe 記憶體頻寬的一半或更差（在 16x 通道 PCIe 3.0 上測試）。

這意味著效能很大程度上取決於具體的使用場景。例如，如果您在每個前向-後向-最佳化器迴圈中驅逐 1 GB 記憶體，那麼在最佳情況下，您可以預期大約 50% 的 PCIe 頻寬作為時間開銷。因此，對於具有 16x 通道的 PCIe 3.0，1 GB 的傳輸速率為 16 GB/s，即每個最佳化器步驟的開銷為 1/(16*0.5) = 1/8 = 125ms。其他開銷可以根據具體的 PCIe 介面、通道以及每次迭代中驅逐的記憶體來估算。

與 CPU 解除安裝相比，如果所有記憶體都適合裝置，分頁最佳化器沒有開銷，只有在需要驅逐部分記憶體時才有一些開銷。對於解除安裝，您通常會解除安裝模型的固定部分，並且每次透過模型迭代時都需要解除安裝和載入所有這些記憶體（有時在前向和後向傳播中各一次）。

< > 在 GitHub 上更新