TorchTPU เป็นโครงการโอเพนซอร์สที่ Google ร่วมมือกับ Meta พัฒนาขึ้น เพื่อมอบแบ็กเอนด์ PyTorch สำหรับ Google TPUs โดยแตกต่างจาก PyTorch/XLA เดิมที่ต้องมีการแก้ไขซอร์สโค้ดอย่างลึกซึ้งเพื่อกำหนดส่วนที่คอมไพล์ได้ ซึ่งเดิมประสบปัญหา Compile-time latency สูง ขาดการรองรับโหมดการทำงานแบบ imperative (“eager”) และการรองรับ Distributed computing primitives ที่เกิดขึ้นใหม่ใน PyTorch ecosystem
TorchTPU มีแนวทางที่แตกต่าง โดยให้ความสำคัญสูงสุดกับการรันแบบ eager execution และช่วยให้เวิร์กโหลด PyTorch เดิมสามารถทำงานได้ด้วยการแก้ไขโค้ดเพียงเล็กน้อย หรือนักพัฒนาสามารถใช้ API ใหม่คือ torch.accelerator.current_accelerator เพื่อระบุ Accelerator ที่ใช้งานอยู่
โหมดการทำงานแบบใหม่นี้เป็นแบบไฮบริดระหว่าง eager/compiled execution ซึ่ง Google อ้างว่าสามารถเร็วกว่า Strict Eager ได้ 50% ขึ้นไป พร้อมลดการใช้หน่วยความจำ HBM สูงสุด นอกจากนี้ Google ยังได้สร้างไปป์ไลน์เพื่อทดสอบโมเดล Hugging Face กว่า 5,000 โมเดล บน TPU รุ่น v5e ขึ้นไป
โครงการนี้มีแผนงานในอนาคตครอบคลุมถึง Dynamic shapes, ประสิทธิภาพการคอมไพล์, การครอบคลุมของโอเปอเรเตอร์, reinforcement learning และ quantization
ศัพท์ที่ควรรู้
- TorchTPU – เป็นโครงการโอเพนซอร์สที่ Google ร่วมกับ Meta พัฒนาขึ้นเพื่อรองรับ PyTorch บน Google TPUs
- PyTorch/XLA – คือเฟรมเวิร์กที่ใช้ในการรันโมเดล ซึ่งเดิมมีข้อจำกัดในการปรับโค้ดและเวลาในการคอมไพล์
- eager execution – คือโหมดการทำงานที่อนุญาตให้โค้ดรันทีละคำสั่ง ทำให้การดีบักง่ายและรวดเร็วกว่าโหมดคอมไพล์
- Distributed computing primitives – คือกลไกพื้นฐานสำหรับการประมวลผลแบบกระจาย ซึ่งเป็นสิ่งที่เพิ่มเข้ามาในระบบนิเวศของ PyTorch
ที่มา: infoq.com