ระบบ Tunix ของ Google ช่วยขจัดจุดคอขวดที่ขัดขวางไม่ให้การเรียนรู้แบบเสริมกำลังเชิงเอเจนต์ (agentic reinforcement learning หรือ RL) ขนาดใหญ่สามารถใช้ TPU ได้อย่างมีประสิทธิภาพ โดยการแยกงานด้านการสร้างข้อมูลปฏิสัมพันธ์ (interaction data) ออกจากงานด้านการอัปเดตนโยบาย (policy) ทำให้ Tunix สามารถผลักดันการใช้งาน TPU จากระดับตัวเลขหลักเดียวขึ้นไปจนเกือบเต็มขีดความสามารถ ซึ่งช่วยลดการสิ้นเปลืองทรัพยากรการคำนวณได้อย่างมหาศาล

จุดคอขวดในการทำ agentic RL

Agentic RL แตกต่างจากการฝึกโมเดลภาษาแบบ "next-token" ที่คุ้นเคยกันดี เนื่องจากเอเจนต์ต้องส่งการเรียกใช้ API, รันโค้ด หรือดำเนินการผ่านสภาพแวดล้อมจำลอง แล้วจึงตอบสนองต่อผลลัพธ์ที่ได้ ดังนั้นลูปการฝึกฝนจึงเป็นแบบซิงโครนัส (synchronous): โมเดลสร้างการกระทำ (action) -> สภาพแวดล้อมทำงาน -> ผลลัพธ์ถูกส่งกลับมา -> และโมเดลจึงจะได้รับเกรเดียนต์อัปเดต (gradient update) เมื่อขั้นตอนในสภาพแวดล้อมเพียงขั้นตอนเดียวต้องใช้เวลาหลายวินาที ฮาร์ดแวร์ TPU ที่มีราคาสูงจึงต้องนั่งว่างงาน และอัตราการใช้งานที่รายงานอาจลดลงต่ำกว่า 10% ซึ่งความไม่มีประสิทธิภาพนี้ส่งผลโดยตรงต่อค่าใช้จ่ายคลาวด์ที่สูงขึ้นและวงจรการวิจัยที่ช้าลง

สถาปัตยกรรมแบบแยกส่วนของ Tunix

Tunix แก้ปัญหานี้โดยการแยกสองขั้นตอน—การสร้างเส้นทาง (trajectory generation) และการเพิ่มประสิทธิภาพนโยบาย (policy optimization)—ออกจากกันไปยังกลุ่มฮาร์ดแวร์ที่แยกกัน

  • Asynchronous actors ทำงานบน CPU หรือ GPU ราคาถูก โดยตัวแสดง (actor) แต่ละตัวจะปฏิสัมพันธ์กับสภาพแวดล้อมที่ได้รับมอบหมายอย่างต่อเนื่อง บันทึกการกระทำและการสังเกตการณ์ แล้วสตรีมเส้นทาง (trajectories) ที่ได้ไปยังพื้นที่จัดเก็บส่วนกลาง
  • Continuous learners ทำงานบน TPU Pods ที่จัดสรรไว้โดยเฉพาะ ตัวเรียนรู้ (learner) จะดึงข้อมูลเป็นชุด (batches) จากบัฟเฟอร์ส่วนกลางและทำการอัปเดตเกรเดียนต์โดยไม่ต้องรอให้ตัวแสดงตัวใดตัวหนึ่งรันโรลเอาต์ (rollout) จนเสร็จสิ้น
  • High-throughput buffer ทำหน้าที่เป็นพื้นที่พักข้อมูลสำหรับเส้นทางต่างๆ เนื่องจากตัวเรียนรู้สามารถอ่านข้อมูลได้เร็วเท่ากับที่บัฟเฟอร์สามารถจ่ายข้อมูลได้ ทำให้ TPU ไม่ต้องหยุดชะงัก

ผลลัพธ์สุทธิคือไปป์ไลน์การฝึกฝนที่ TPU ทำงานหนักเกือบตลอดเวลา โดยผลักดันการใช้งานให้เข้าใกล้ 100%

อุปสรรคทางเทคนิคและวิธีที่ Tunix เอาชนะ

เอพิโซดที่มีความยาวไม่คงที่และการคอมไพล์ซ้ำของ XLA

คอมไพเลอร์ XLA ของ JAX จะปรับแต่งเพื่อรูปทรงเทนเซอร์แบบคงที่ (fixed tensor shapes) อย่างไรก็ตาม งานเชิงเอเจนต์มักจะสร้างลำดับที่มีความยาวแตกต่างกัน ซึ่งตามปกติจะทำให้เกิดการคอมไพล์ซ้ำ (recompilation) ที่สิ้นเปลือง Tunix จึงใช้วิธีแพ็กลำดับที่สั้นกว่าเข้าด้วยกัน และจัดกลุ่มเอพิโซดที่มีความยาวใกล้เคียงกันลงในถัง (buckets) เพื่อรักษาความเสถียรของรูปทรงข้อมูลให้นานพอที่ XLA จะนำเคอร์เนล (kernels) ที่คอมไพล์แล้วมาใช้ซ้ำได้ ผลลัพธ์ที่ได้คือปริมาณงาน (throughput) ที่สม่ำเสมอโดยไม่มีภาระส่วนเกินจากคอมไพเลอร์ที่จะมาบั่นทอนประสิทธิภาพ

การขยายขนาดโมเดลขนาดมหึมาผ่านชิป TPU จำนวนมาก

การฝึกเอเจนต์ที่มีพารามิเตอร์มากกว่า 70 พันล้านตัว จำเป็นต้องกระจายน้ำหนัก (weights) และข้อมูลไปยังโหนด TPU หลายโหนด Tunix ใช้พรอมิทิฟ ShardMap ของ JAX เพื่อทำ sharding ทั้งพารามิเตอร์ของโมเดลและการกระตุ้น (activations) ช่วยให้ตัวเรียนรู้สามารถเก็บโมเดลทั้งหมดไว้ในหน่วยความจำในขณะที่ยังสามารถป้อนข้อมูลด้วยความเร็วสูงได้ กลยุทธ์การทำ sharding นี้ทำให้สามารถฝึกโมเดลที่ก่อนหน้านี้ไม่สามารถทำได้ด้วย TPU pod เพียงชุดเดียว

เกรเดียนต์ที่ล้าสมัยจากไปป์ไลน์แบบแยกส่วน

เมื่อตัวแสดงทำงานล่วงหน้าตัวเรียนรู้ ข้อมูลที่พวกเขาส่งมาอาจกลายเป็นข้อมูลที่ "ล้าสมัย" (stale) เมื่อเทียบกับนโยบายปัจจุบัน Tunix บรรเทาปัญหานี้ด้วยสองกลไก: การใช้ importance-sampling เพื่อถ่วงน้ำหนักตัวอย่างเก่าใหม่เพื่อให้สะท้อนถึงความเกี่ยวข้อง และการกำหนดเกณฑ์ความล้าสมัย (staleness threshold) เพื่อทิ้งเส้นทางที่มีอายุเกินกว่าที่กำหนดไว้ ทั้งสองส่วนนี้ช่วยให้การเรียนรู้มีความเสถียรแม้ว่าไปป์ไลน์จะทำงานแบบอะซิงโครนัสก็ตาม

สิ่งที่ผู้ใช้งานต้องระวัง

  • Latency audit – ประโยชน์ของการแยกส่วนขึ้นอยู่กับเวลาตอบสนองของสภาพแวดล้อม ทีมงานควรวัดความหน่วงแบบ end-to-end และตรวจสอบให้แน่ใจว่ากลุ่มตัวแสดงมีขนาดที่เหมาะสมเพื่อให้บัฟเฟอร์มีข้อมูลเต็มอยู่เสมอ
  • Worker pool design – CPU หรือ GPU ราคาถูกสามารถรองรับตัวแสดงได้จำนวนมาก แต่การใช้งานเกินขีดจำกัด (oversubscribing) อาจทำให้เกิดการแย่งชิงทรัพยากรบนเครือข่ายหรือพื้นที่จัดเก็บข้อมูล การออกแบบกลุ่มผู้ปฏิบัติงานที่สมดุลและสอดคล้องกับอัตราการรับข้อมูลของบัฟเฟอร์จึงเป็นสิ่งสำคัญ
  • Buffer robustness – ระบบจัดเก็บส่วนกลางต้องรองรับอัตราการเขียนและอ่านที่สูงโดยไม่กลายเป็นจุดคอขวดใหม่ การเลือกใช้ระบบจัดเก็บข้อมูลที่มีค่า tail latency ต่ำและมีแบนด์วิดท์เพียงพอถือเป็นส่วนประกอบที่ขาดไม่ได้ของสถาปัตยกรรมนี้

ข้อเสียที่อาจเกิดขึ้น

สถาปัตยกรรมแบบแยกส่วนทำให้มีส่วนประกอบที่ต้องจัดการมากขึ้น ทั้งกลุ่มฮาร์ดแวร์ที่แยกจากกัน, บัฟเฟอร์ที่ต้องทำงานตลอดเวลา และตรรกะการประสานงานเพื่อควบคุมขีดจำกัดความล้าสมัย

บทสรุป

Tunix แสดงให้เห็นว่าต้นทุนหลักในการทำ agentic RL ไม่ใช่ตัวโมเดลเอง แต่เป็นเวลาว่างงานที่เกิดจากลูปการปฏิสัมพันธ์แบบซิงโครนัส การย้ายงานโรลเอาต์ไปยังฮาร์ดแวร์ราคาถูก และการป้อนข้อมูลเข้าสู่ TPU pod ที่เรียนรู้อย่างต่อเนื่องจากบัฟเฟอร์ที่มีความสามารถในการรับส่งข้อมูลสูง ทำให้ Google สามารถเปลี่ยนปัญหาการใช้งานที่ต่ำกว่า 10% ให้กลายเป็นเวิร์กโฟลว์ที่ทำงานได้เกือบเต็มขีดความสามารถ