機器學習領域的工程師們對於 PyTorch/XLA 一直抱有很高的期待。這款 Python 套件通過 XLA 深度學習編譯器連接 PyTorch 框架和 Cloud TPU，帶來更強大的運算能力。現在，PyTorch/XLA 推出 2.5 版本，不僅增添了對 vLLM 的支援，還在開發者體驗方面進行了多項優化，讓使用者能以更簡單的方式享受深度學習在 Cloud TPU 上的強大性能。
文章重點
vLLM 支援拓展到 TPU，提升模型服務效能
vLLM 是一款設計精良的高吞吐量推理引擎，以其高效的記憶體管理和 Hugging Face Model Hub 直接集成而聞名。PyTorch/XLA 2.5 的一大亮點在於提供了實驗性的 vLLM TPU 支援，讓開發者能夠在 TPU 上無縫運行 vLLM。這意味著只需更改少量配置，vLLM 使用者便能將後端從 GPU 切換至 TPU，並享受相同的服務介面和高效能推理體驗。Pallas 核心功能，包括 paged attention 和 flash attention，以及效能優化的 dynamo bridge，現已支援 TPU，並納入了 PyTorch/XLA 代碼庫，為 TPU 用戶帶來了更多的選擇。
精簡 API，降低開發者學習成本
此次更新的 PyTorch/XLA 2.5 還加強了與原生 PyTorch API 的一致性，將部分自定義 API 遷移到 PyTorch 本身，使開發者更容易上手。例如，分布式 API 函數現在已從 torch_xla 遷移到 torch.distributed，降低了新使用者的學習門檻，同時優化了使用體驗。未來還將有更多的 API 合併至 PyTorch，以進一步簡化開發流程。
更強大的 torch_xla.compile 調試功能
在開發與調試模型時，開發者需要高效的工具來迅速發現問題。PyTorch/XLA 2.5 對 torch_xla.compile 函數進行了一系列改進，讓調試更加靈活。新增的「full_graph」模式可以在模型出現多個編譯圖形時發出警告，便於提前發現潛在問題。此外，開發者還可以設定編譯函數的重編譯次數，避免因模型的意外動態行為而導致過度編譯。針對命名功能的改進則讓調試信息更具可讀性，讓開發者能輕鬆定位問題。
開始使用 PyTorch/XLA 2.5
對於機器學習工程師來說，PyTorch/XLA 2.5 為 Cloud TPU 的開發帶來了更多便利與靈活性。無論是希望提升推理效能的 vLLM 使用者，還是尋求一致 API 和高效調試體驗的開發者，都能在此次更新中找到所需的功能。立即通過 Python 包管理器下載最新版的 PyTorch/XLA 2.5，或瀏覽 GitHub 項目頁了解更多，體驗新版所帶來的性能和效率提升！
