กูเกิลเผยรายละเอียดของซูเปอร์คอมพิวเตอร์ที่ใช้ชิปออกแบบเอง Tensor Processing Unit (TPU) v4 ซึ่งเริ่มใช้ในโปรดักชันมาตั้งแต่ปี 2020 (แต่เพิ่งเผยรายละเอียดปี 2023) ว่าสามารถยกระดับประสิทธิภาพ machine learning ได้เกือบ 10 เท่าจากเครื่อง TPU v3 และสามารถเอาชนะเครื่องที่ใช้จีพียู NVIDIA A100 ได้ด้วย

TPU v4 เปิดตัวต่อสาธารณะเมื่อปี 2021 และทำผลงานเบนช์มาร์คด้าน AI ได้ดี เรื่องใหม่ที่กูเกิลเปิดเผยเพิ่มเติมในรอบนี้คือเครื่อง TPU v4 มีฟีเจอร์สำคัญ 2 ประการ

ใช้สวิตช์แบบแสง (Optical Circuit Switching)
มีตัวช่วยประมวลผลข้อมูลชื่อ SparseCore

ฟีเจอร์ 2 อย่างนี้ช่วยให้เครื่อง TPU v4 สามารถรีดประสิทธิภาพงานด้าน AI ได้ดีขึ้นกว่าเดิม

Optical Circuit Switching

ชิป TPU ถูกออกแบบมาให้ต่อกันเป็นคลัสเตอร์ขนาดใหญ่ มีชิปจำนวนมาก ซึ่งกูเกิลตั้งเป้าว่า TPU v4 ต้องมีจำนวนชิปเพิ่มขึ้น 4 เท่าจาก TPU v3 (4096 ตัว vs 1024 ตัว) จึงเริ่มเจอข้อจำกัดในการเชื่อมต่อ (interconnect) ระหว่างชิปแต่ละตัว เพราะระยะห่างระหว่างชิป (ที่อาจต้องอยู่ข้ามแร็คกัน) เริ่มไกลเกินการส่งข้อมูลด้วยสัญญาณไฟฟ้า ทางออกเดียวจึงเป็นการใช้แสง (optical link)

ทีมเครือข่ายของกูเกิลได้พัฒนาสวิตช์แสง (Optical Circuit Switching) ชื่อว่า Palomar โดยอาศัยเทคโนโลยีกระจกแบบ 3D Micro-Electro-Mechanical Systems (MEMS) ช่วยสลับวงจรได้รวดเร็วระดับมิลลิวินาที รายละเอียดอ่านได้จากบล็อก Jupiter evolving: Reflecting on Google’s data center network transformation

สวิตช์แบบ OCS ยังทำใช้เชื่อมต่อชิป TPU v4 แบบ 3D ได้ (ของ TPU v3 ต่อเป็น 2D) ช่วยให้แบนด์วิดท์ระหว่างกันเพิ่มขึ้น และลดปัญหาชิปบางตัวไม่ทำงานได้ด้วย ด้วยข้อจำกัดของพื้นที่แร็คทำให้กูเกิลเลือกต่อเป็น 4x4x4 ชุดละ 64 ตัว

ตัวอย่างเครื่อง TPU v4 ของจริง จำนวน 8 ชุด ซึ่งคิดเป็น 1/8 ของทั้งระบบ

TPU v4 ใช้สวิตช์แสง OCS แบบ reconfigurable คือปรับแต่งการวาง topology ของเครือข่ายได้ตามต้องการ ทำให้โครงสร้างของการวางชิป TPU v4 ยืดหยุ่นตามเวิร์คโหลดของงาน AI ประเภทต่างๆ ได้

ผลในภาพรวมคือ สวิตช์แบบ OCS ทำงานได้เร็วกว่า ประหยัดพลังงานกว่า ราคาถูกกว่า การเชื่อมต่อแบบ Infiniband (ปัจจุบันเป็นของ NVIDIA) ที่นิยมใช้ในเซิร์ฟเวอร์สมรรถนะสูง ตอนนี้ชิ้นส่วน OCS มีต้นทุนน้อยกว่า 5% ของระบบ TPU v4 ทั้งหมด และใช้พลังงานน้อยกว่า 5% ของทั้งระบบเช่นกัน

ตัวช่วยประมวลผล SparseCore

งานประมวลผล AI มีหลายประเภท โดยงานประเภทหนึ่งที่พบบ่อยคือ Deep learning recommendation models (DLRM) ซึ่งเป็นงานแนะนำสิ่งที่ผู้ใช้น่าจะชอบ ที่เราคุ้นเคยกันจากผลิตภัณฑ์ต่างๆ ของกูเกิล เช่น Search, Ads, YouTube, Google Play

อัลกอริทึมแบบ DLRM จะมีชั้นของการนำข้อมูลที่กระจัดกระจายมาจัดให้เป็นหมวดหมู่ เรียกว่า embedding ซึ่งจะช่วยให้เทรนโมเดลได้เร็วขึ้น อย่างไรก็ตาม ตารางข้อมูลที่ใช้เปรียบเทียบเพื่อทำ embedding มักมีขนาดใหญ่ ใช้แรมเยอะ แต่ประมวลผลจริงๆ น้อย (รูปแบบเวิร์คโหลดคือหาข้อมูลในตารางขนาดใหญ่) จึงกลายเป็นคอขวดของแรม

เมื่องานประเภท DLRM มีสัดส่วนการใช้งานราว 25% ของเวิร์คโหลดทั้งหมดในระบบ กูเกิลจึงประดิษฐ์ชิปชื่อ SparseCore มาประมวลผลงานส่วน embedding โดยเฉพาะ แยกจากชิป TensorCore ที่ใช้ประมวลผลโมเดล AI หลัก ซึ่งมีธรรมชาติแตกต่างกัน

ชิป SparseCore เอาเข้าจริงแล้วเริ่มใช้งานมาตั้งแต่ TPU v2 (แต่เพิ่งมาเผยตัวต่อชาวโลก) และปรับปรุงมาเรื่อยๆ ใน TPU v3 และ v4 จนตอนนี้ช่วยเพิ่มประสิทธิภาพของเวิร์คโหลด DLRM ได้มาก โดยใช้พลังงานเพียง 5% และพื้นที่ชิปเพียง 5% ของ TPU ทั้งหมดเท่านั้น

ประสิทธิภาพของ TPU v4 เพิ่มขึ้นราว 3 เท่าจาก TPU v3 และหากเทียบกับการใช้ซีพียูมาตรฐานก็สูงกว่ากันถึง 30 เท่าเลยทีเดียว

ด้วยเทคนิคการปรับแต่งประสิทธิภาพของ TPU v4 หลายอย่างข้างต้น ทำให้เครื่อง TPU v4 สามารถเอาชนะ TPU v3 ได้สบายๆ ประสิทธิภาพต่อวัตต์ดีขึ้น 40%, ประสิทธิภาพดีขึ้นราว 1.5-3 เท่า

กูเกิลยังนำเครื่อง TPU v4 ไปเทียบกับเครื่องที่ใช้ชิปประมวลผล AI ของคู่แข่งคือ NVIDIA A100 (ยังไม่ใช่ H100 ตัวใหม่ล่าสุดที่เพิ่งเริ่มออกขายปีนี้) และ Graphcore MK2 IPU ก็สามารถรันเบนช์มาร์ค MLPerf 2.0 เอาชนะ A100 ได้เช่นกัน

ตัวอย่างลูกค้าที่รันงานบน TPU v4 ผ่าน Google Cloud คือ Midjourney บริการ AI สร้างรูปจากข้อความชื่อดัง และสถาบัน Allen Institute for AI ที่ก่อตั้งโดย Paul Allen ผู้ร่วมก่อตั้งไมโครซอฟท์

ที่มา – Google Cloud Blog, เปเปอร์ฉบับเต็ม (PDF), The Register