不依賴CUDA的大模型推理已經(jīng)實現(xiàn) 原創(chuàng)
在過去的十年中，CUDA已經(jīng)成為深度學(xué)習(xí)模型推理領(lǐng)域的核心技術(shù)。NVIDIA的CUDA技術(shù)優(yōu)化了GPU的性能，提高了數(shù)據(jù)處理的速度和效率。然而，CUDA的主導(dǎo)地位也帶來了對NVIDIA硬件的依賴，這限制了在非NVIDIA平臺上進行高效推理的可能性。為了應(yīng)對這一挑戰(zhàn)，近年來，開源社區(qū)和企業(yè)開始尋求開發(fā)不依賴于特定供應(yīng)商的推理技術(shù)。
Triton的崛起
OpenAI開發(fā)的Triton編程語言應(yīng)運而生，旨在為廣泛的硬件設(shè)備提供高效的并行編程方式。
它允許開發(fā)者編寫一次代碼，就能在NVIDIA、AMD甚至未來的Intel GPU上運行，極大地擴展了AI模型的適用范圍。Triton的出現(xiàn)，不僅打破了CUDA的壟斷，還為AI社區(qū)帶來了新的活力和創(chuàng)新的可能性。
與CUDA相比，Triton提供了一個更靈活、更高級的抽象層，使得開發(fā)者可以在多種類型的GPU上實現(xiàn)優(yōu)化的計算性能，而不僅僅是NVIDIA的產(chǎn)品。
Triton的核心優(yōu)勢在于其對于大型模型如LLM的支持。Triton能夠處理諸如矩陣乘法、正規(guī)化（RMSNorm）、自注意力機制（RoPE）、激活函數(shù)（如SiLU）以及元素級乘法等關(guān)鍵操作，這些都是LLM推理中不可或缺的部分。PyTorch社區(qū)已經(jīng)利用Triton重寫了這些操作的內(nèi)核，通過torch.compile實現(xiàn)了自動生成Triton內(nèi)核的能力。
技術(shù)實現(xiàn)細節(jié)
將CUDA內(nèi)核替換為Triton內(nèi)核是一個復(fù)雜的過程，需要精細的調(diào)優(yōu)和定制。例如，PyTorch團隊開發(fā)了一種自定義的FP16 Triton GEMM（General Matrix-Matrix Multiply）內(nèi)核，該內(nèi)核專門針對LLM中的線性層的矩陣乘法進行優(yōu)化。這個過程涉及到對不同線性層權(quán)重矩陣形狀的精確調(diào)整，以確保最佳性能。
此外，flash attention是LLM中另一個關(guān)鍵操作，PyTorch團隊采用了多種配置的Triton flash attention內(nèi)核進行實驗，最終選擇了AMD的Triton flash attention內(nèi)核，因為它能夠在eager和compile模式下提供清晰的輸出，同時兼容torch.compile。
性能評估
在性能方面，雖然Triton內(nèi)核的性能尚未完全達到CUDA的水平，但已經(jīng)非常接近。在NVIDIA的H100和A100 GPU上的測試顯示，使用Triton內(nèi)核的模型推理性能達到了CUDA模型性能的78%到82%。這一結(jié)果足以證明，即使是在高要求的應(yīng)用場景下，非CUDA推理也能達到實用的水平。
未來發(fā)展方向
展望未來，Triton內(nèi)核還有很大的優(yōu)化空間。PyTorch團隊和OpenAI已經(jīng)在探索更高效的內(nèi)核設(shè)計，如利用更精細的工作分解策略和持久內(nèi)核技術(shù)，這可能會進一步提高Triton的計算效率并縮小與CUDA的性能差距。此外，隨著FlexAttention和FlashAttention-3等新技術(shù)的引入，PyTorch有望在更廣泛的硬件平臺上實現(xiàn)更長的上下文長度和更大規(guī)模的解碼任務(wù)，這將推動整個行業(yè)向前發(fā)展。
Triton中國社區(qū)的發(fā)展和影響
2024年，智源研究院舉辦了多期Triton中國生態(tài)Meetup活動，這些活動不僅聚焦于技術(shù)分享和經(jīng)驗交流，也致力于構(gòu)建一個開放、協(xié)作的Triton技術(shù)生態(tài)體系。國內(nèi)大量支持Triton語言編程的芯片廠商積極參與了交流。
隨著Triton和其他非CUDA技術(shù)的成熟，社區(qū)正逐步克服對特定硬件供應(yīng)商的依賴，這將使得大模型推理技術(shù)更加普及和可用。這一進展不僅對科研人員和開發(fā)者具有重要意義，也對推動AI技術(shù)的普及具有深遠影響。隨著更多的創(chuàng)新和研究，未來的大模型推理將更加高效、靈活和開放。
參考鏈接：
1. https://pytorch.org/blog/cuda-free-inference-for-llms/
2. https://hub.baai.ac.cn/view/39486
本文轉(zhuǎn)載自公眾號AI時代窗口 作者：郁愈
