HomebrewML/HeavyBall

What it solves

HeavyBall 提供了一個為 PyTorch 設計的高性能、可組合的優化器函式庫。它透過使用可編譯的構建模組並將其融合進 Triton kernels，解決了標準優化器實作中的低效率問題，從而減少記憶體流量並提高執行速度。它還簡化了進階優化技術（例如二階方法和狀態壓縮）的整合，這些技術通常難以實作或組合。

How it works

該函式庫由超過 100 個編譯後的函數組成，並組裝成轉換鏈。透過使用 torch.compile(fullgraph=True)，這些函數會融合進極簡的 kernels，顯著減少對記憶體的讀寫次數。例如，一個標準的 Adam 更新可以從 14 次讀取和 9 次寫入減少到 4 次讀取和 3 次寫入。

Who it’s for

機器學習工程師和研究人員，他們需要更快的優化器步驟、更低的優化器狀態記憶體開銷，或者需要在 PyTorch 生態系統中存取各種進階的二階和正交優化演算法。

Highlights

• Extensive Optimizer Suite: 包括與 AdamW、SGD 和 RMSprop API 相容的替換方案，以及 Muon、SOAP、LATHER 和 ADOPT。
• Composable Features: 可鏈接的標記，用於 MARS 方差減少、謹慎更新以及 PaLM 風格的 beta2 調度。
• ECC State Compression: 減少優化器狀態的記憶體使用量（例如，bf16 + int8 修正），以在不犧牲精度的情況下節省空間。
• High Performance: 透過 torch.compile 將操作融合進 Triton kernels，從而顯著提升步驟延遲的執行速度。
• Distributed Support: 原生相容於 DDP 和 FSDP，包括針對複雜二階方法的自動重新分區。