PyTorch釋出可大幅改進AI模型運算效能的torchao函式庫

PyTorch推出全新原生函式庫torchao，透過低精度資料型態、量化和稀疏性技術，減少模型的計算成本和記憶體使用量，讓模型執行更高效能。torchao提供了一套易於使用的工具組，支援多種模型推理和訓練最佳化方法，可廣泛用於PyTorch模型，熱門的LLaMA 3和Diffusion模型皆可獲得顯著的效能提升。

低精度資料型態是torchao加速的關鍵之一。該工具支援float8、int4等低精度資料型態，能有效減少計算成本和記憶體使用需求，像是在LLaMA 3 70B模型預訓練中，torchao提供了float8訓練流程，可將模型運算加速達1.5倍。開發者只需要使用convert_to_float8_training函式，即可將模型訓練運算轉換成float8，簡單實現高效訓練。

在推理方面，torchao提供多種量化方式，包括權重量化（Weight-Only Quantization）和動態啟動量化（Dynamic Activation Quantization）。使用者可以透過API自行選擇適合的量化策略，以達到推理的最佳效能。

torchao提供int4權重量化方法，可以將LLaMA 3 8B模型的推理速度提升達97%，而透過量化的鍵值快取，在128K上下文長度下顯示記憶體使用量可減少73%。此外，torchao的autoquant功能可自動替模型的每一層選擇最佳量化方式，減少開發者手動調整的麻煩。

稀疏性技術也是torchao提升模型效率的重要手段，藉由引入稀疏性，torchao可最佳化模型參數計算，降低不必要的運算開銷，例如torchao結合稀疏性和量化技術，可讓ViT-H模型推理獲得5％速度提升。同時，torchao還可將權重量化成int4，並將鍵值快取量化成int8，使LLaMA 3.1 8B在完整的128K上下文長度下，僅佔用18.9GB的顯示記憶體。

torchao不僅顯著提升LLaMA 3 8B模型的基準測試效能，還讓Diffusion模型Flux.1-Dev的推理速度提升53.88％。

PyTorch官方未來將繼續更新torchao，支援精度低於int4以下的資料型態，開發更高吞吐量的推理核心，甚至拓展支援更多的硬體後端。