【ITBEAR】PyTorch官方近日宣布，旗下全新架構(gòu)優(yōu)化庫(kù)torchao已正式上線。該庫(kù)專注于模型的量化和稀疏性優(yōu)化，致力于在保障性能的同時(shí)，降低AI模型的計(jì)算成本及RAM占用，從而提升其運(yùn)行效率。
torchao提供了一整套優(yōu)化工具，特別針對(duì)如LLaMA 3等熱門AI模型進(jìn)行性能提升。通過(guò)支持float8、int4等低精度數(shù)據(jù)類型，該庫(kù)能有效減輕硬件負(fù)擔(dān)，減少內(nèi)存使用。
在LLaMA 3 70B模型的預(yù)訓(xùn)練過(guò)程中，torchao展現(xiàn)了顯著效果。其float8訓(xùn)練流程能提升模型計(jì)算速度達(dá)1.5倍，開(kāi)發(fā)者僅需調(diào)用convert_to_float8_training函數(shù)，便可輕松實(shí)現(xiàn)模型的高效訓(xùn)練轉(zhuǎn)換。
在推理方面，torchao提供了包括權(quán)重量化和動(dòng)態(tài)激活量化在內(nèi)的多種量化策略，供用戶靈活選擇，以達(dá)到最佳的模型推理效果。
在稀疏性優(yōu)化領(lǐng)域，torchao同樣表現(xiàn)出色。它不僅能提升模型參數(shù)的計(jì)算效率，還能顯著降低顯存占用。例如，在ViT-H模型的推理中，torchao可提升速度達(dá)5%。同時(shí)，通過(guò)將LLaMA 3.1 8B模型的權(quán)重量化為int4，并將鍵值緩存量化為int8，torchao使得該模型在完整的128K上下文長(zhǎng)度下，顯存占用僅需18.9GB。
