gomlx/gomlx

解決的問題

GoMLX 為 Go 程式語言提供一個全面的機器學習與數學框架，作為 PyTorch、JAX 或 TensorFlow 等流行框架的 Go 原生替代方案。它讓開發者能在 Go 生態系內訓練、微調與執行機器學習模型，減少對 Python 進行 ML 開發與生產的依賴。

工作原理

GoMLX 作為一個中介 API，提供一系列 ML 層與工具庫，並透過可交換的後端執行：

• XLA 後端：使用 OpenXLA 實現高效率 JIT 編譯，支援 CPU、GPU（Nvidia、AMD ROCm、Intel、Mac）與 TPU。支援跨多個加速器的分散式執行。
• Go 後端：純 Go 實作，高度可移植，支援 WASM 用於瀏覽器執行，並針對特定架構使用 SIMD 優化。
• ONNX 後端：利用 ONNX Runtime 執行運算，並支援將模型儲存為 .onnx 格式。

該框架包含自動微分（autodiff）、變數管理（Store 與 Scope）系統，以及大量可微分運算子與 ML 層。

適用對象

專為希望在不離開 Go 生態系的情況下建構、實驗並部署機器學習模型的 Go 開發者設計，也適用於探索非傳統最佳化器或正則化方法的研究人員。

主要亮點

• HuggingFace 集成：與 go-huggingface 相容，支援下載模型、分詞器、資料集，並從 GGUF 或 safetensors 讀取參數。
• 多樣化的層支援：包含 FFN、卷積、多頭注意力、LSTM，以及 KAN（Kolmogorov-Arnold 網路）與 VNN（向量神經網路）等特殊層。
• 彈性執行：透過 XLA 支援多種硬體目標，包括多 GPU/TPU 分散式執行。
• 開發者工具：包含用於檢查檢查點與產生圖表的命令列工具，並透過 GoNB 與 JupyterLab 集成。
• 進階訓練功能：支援梯度檢查點，可在大型模型訓練時以記憶體換取重計算。