เร่งการอนุมาน AI สำหรับการสร้าง 3 มิติบน Roblox

• Roblox ได้นำ CUDA Graphs และ KV caching มาใช้เพื่อเพิ่มความเร็วในการสร้างตาข่าย 3 มิติสำหรับการทำงานซ้ำที่ตอบสนองได้ดียิ่งขึ้น
• เมื่อเปิดตัว โมเดล Cube 3D สามารถสร้างโทเค็นได้ในเวลา 7.8 มิลลิวินาที (ลดลงจาก 60.5 มิลลิวินาที) และสร้างวัตถุเต็มรูปแบบได้ในเวลา 4 วินาที (ลดลงจาก 31 วินาที) 

เมื่อต้นปีนี้ Roblox ได้แบ่งปันความสามารถแรกจากโมเดลพื้นฐาน Cube 3D ของเรา ด้วย Cube 3D ผู้สร้างสามารถสร้างโมเดล 3 มิติและสภาพแวดล้อมได้โดยตรงจากข้อความที่ป้อน ตั้งแต่เริ่มต้น เราให้ความสำคัญกับการเพิ่มประสิทธิภาพความหน่วงเวลา โดยตระหนักดีว่าเวลาการสร้างที่ช้าจะรบกวนกระบวนการที่โดยธรรมชาติแล้วต้องทำซ้ำ ก่อนที่ Cube 3D จะเปิดตัวในเดือนมีนาคม เราได้ทำให้ขั้นตอนการอนุมานเร็วขึ้น 7.8 เท่าและตอบสนองได้ดีขึ้นทั้งสำหรับนักพัฒนาและผู้ใช้แล้ว 

ตั้งแต่เปิดตัว มีวัตถุมากกว่า 578,000 ชิ้นถูกสร้างขึ้นในประสบการณ์ที่น่าสนใจหลายแห่ง นักพัฒนายังแสดงความสนใจในการเปิดโอกาสให้ผู้ใช้สามารถสร้างวัตถุ 3 มิติภายในประสบการณ์ต่างๆ ได้ด้วยข้อความคำสั่ง เช่น "แมว", "เบอร์เกอร์" เป็นต้น ที่โดดเด่นที่สุดคือ Mic Up เกมแฮงเอาท์ยอดนิยมที่ใช้การแชทด้วยเสียง ได้ใช้ Cube 3D เพื่อมอบวิธีที่สนุกและมีปฏิสัมพันธ์ให้กับผู้เล่นในการสร้างวัตถุ ในการใช้งาน ผู้เล่นสามารถเปิดเมนูด้านซ้ายเพื่อดูความสามารถเพิ่มเติม รวมถึงไอคอน AI หลังจากคลิกที่ไอคอนนั้น ผู้เล่นสามารถป้อนข้อความเพื่อสร้างวัตถุ 3 มิติได้ สำหรับผู้ใช้ เวลาในการสร้างที่ยาวนานขึ้นจะสร้างความไม่สะดวก ทำให้พวกเขาพลาดความมหัศจรรย์ของการเห็นความคิดของพวกเขาถูกแปลงเป็นวัตถุ 3 มิติแบบเรียลไทม์ 

เราต้องการเปลี่ยนประสบการณ์ของผู้ใช้ในยุค 3D จากการโต้ตอบแบบหยุดและรอ ให้กลายเป็นสิ่งที่ตอบสนองและรู้สึกเป็นธรรมชาติ ซึ่งช่วยให้สามารถทดลองได้อย่างรวดเร็ว ความสามารถในการเพิ่มวัตถุลงในฉากได้อย่างรวดเร็วนั้นเป็นสิ่งสำคัญสำหรับนักพัฒนา เพื่อเร่งความเร็วของ Cube 3D เราได้ทำการวิเคราะห์โปรไฟล์ของกระบวนการอนุมานก่อนเป็นอันดับแรก เพื่อระบุจุดคอขวดด้านประสิทธิภาพ แม้ว่าจะใช้ GPU ที่ทรงพลัง แต่เราก็พบว่ามีเวลาว่างระหว่างการทำงานอยู่มาก

กรอบการทำงานการเรียนรู้เชิงลึกสมัยใหม่พึ่งพา CPU ในการจัดตารางและเริ่มดำเนินการ (หรือเคอร์เนล) บน GPU CPU เตรียมแต่ละการดำเนินการ ส่งไปยัง GPU และรอการยืนยันก่อนที่จะเตรียมการดำเนินการถัดไป การรอคอยนี้สร้างคอขวดในการจัดตารางที่ GPU อาจนั่งว่างในขณะที่ CPU เตรียมชุดงานถัดไป ในอุดมคติ เราต้องการให้ CPU ทำงานนำหน้า GPU เตรียมและคิวการดำเนินการเพื่อให้ GPU มีงานทำอยู่เสมอ

นี่เป็นปัญหาโดยเฉพาะอย่างยิ่งสำหรับตัวถอดรหัสแบบออโตเรเกรสซีฟในโมเดลประเภททรานส์ฟอร์เมอร์ เช่น Cube 3D ซึ่งจำเป็นต้องประมวลผลอินพุตและสร้างโทเค็นทีละลำดับ โมเดลเหล่านี้ต้องการการดำเนินการเป็นรายบุคคลนับพันครั้งสำหรับการสร้างเพียงหนึ่งครั้ง และภาระการคำนวณจะสะสมเพิ่มขึ้นในแต่ละขั้นตอนของลำดับ

"เราต้องการสร้างสิ่งที่ช่วยให้เกิดการโต้ตอบแบบสี่มิติ" รองประธานฝ่ายวิศวกรรม อนุปัม ซิงห์ กล่าวอธิบายเหตุผลที่ Roblox เลือกใช้แนวทางออโตเรเกรสซีฟ "เราไม่ได้ต้องการแค่สร้างรถยนต์เท่านั้น แต่เรายังต้องการที่จะเปิดประตูรถและเข้าไปข้างในได้ด้วย"

แต่ละการดำเนินการมีค่าใช้จ่าย:

• เวลา CPU ที่ใช้ในการเตรียมแต่ละเคอร์เนล
• ค่าใช้จ่ายที่เกิดขึ้นจากการเปิดตัวเคอร์เนล
• เวลาการประมวลผลของ GPU (การคำนวณจริง)
• ค่าใช้จ่ายในการซิงโครไนซ์เมื่อตรวจสอบการเสร็จสิ้น

ในกรณีของการดำเนินงานขนาดเล็กที่ทำงานได้อย่างรวดเร็วบน GPU ค่าใช้จ่ายเพิ่มเติมนี้อาจครอบงำเวลาในการอนุมาน GPU อาจทำการคำนวณอย่างจริงจังเพียงเศษส่วนเล็กน้อยของเวลาในการอนุมานทั้งหมด

เพื่อแก้ไขปัญหาคอขวดนี้ เราได้ใช้ประโยชน์จาก CUDA Graphs ซึ่งเป็นฟีเจอร์ที่ช่วยให้สามารถบันทึกและเล่นลำดับของการดำเนินการบน GPU ได้โดยไม่ต้องใช้ CPU เข้ามาแทรกแซง ส่วนประกอบตัวถอดรหัสแบบออโตเรเกรสซีฟของสถาปัตยกรรม Cube 3D จะประมวลผลข้อความคำสั่งและสร้างโทเค็นรูปทรงผ่านเวกเตอร์ที่มีความยาวคงที่ 

แม้ว่าสถาปัตยกรรมตัวถอดรหัสแบบสองกระแสของเราจะมีความคล้ายคลึงกับโมเดลภาษาขนาดใหญ่แบบดั้งเดิม (LLM) ในแง่ของการทำงาน แต่มีความแตกต่างที่สำคัญ—คือมันใช้กระแสความสนใจแบบคู่ขนาน สองกระแส กระแสหนึ่งใช้สำหรับโทเค็นเงื่อนไข และอีกกระแสหนึ่งสำหรับโทเค็นรูปร่าง เครื่องมือการอนุมาน LLM ที่มีอยู่ในตลาดไม่เหมาะกับความต้องการของเรา และเราจำเป็นต้องมีการพัฒนาเฉพาะที่ปรับให้เข้ากับสถาปัตยกรรมของเราโดยเฉพาะ

คิดถึงกราฟ CUDA เหมือนกับการบันทึกมาโครสำหรับ GPU แทนที่ CPU จะออกคำสั่งแต่ละคำสั่งเป็นรายบุคคล มันจะบันทึกลำดับการดำเนินการของ GPU ทั้งหมด (กราฟ) และเปิดตัวกราฟทั้งหมดด้วยคำสั่ง CPU เพียงคำสั่งเดียว วิธีการนี้ช่วยลดภาระการเปิดตัวเคอร์เนลได้อย่างมากโดยการกำจัดความจำเป็นที่ CPU จะต้องกำหนดเวลาการดำเนินการแต่ละรายการในระหว่างการอนุมาน เมื่อกราฟถูกเปิดตัวแล้ว GPU จะดำเนินการลำดับทั้งหมดโดยอัตโนมัติโดยไม่ต้องรอคำสั่งเพิ่มเติม

กราฟ CUDA มีข้อจำกัดบางประการ เนื่องจากโครงสร้างของกราฟจำเป็นต้องถูกกำหนดไว้ล่วงหน้า จึงต้องใช้ขนาดแบตช์และมิติของอินพุตที่คงที่ ซึ่งหมายความว่าจำเป็นต้องสร้างกราฟแยกสำหรับแต่ละขนาดแบตช์หรือรูปร่างของอินพุต สำหรับกรณีการใช้งานของเราด้วย Cube 3D ข้อจำกัดนี้ถือว่ายอมรับได้ เนื่องจากเราสามารถทำให้กระบวนการอนุมานเป็นมาตรฐานโดยอิงตามมิติของอินพุตที่พบบ่อย

เราจำเป็นต้องปรับแนวทางของเราเพื่อนำ CUDA Graphs มาใช้กับโมเดล Cube 3D ของเรา ใน LLM แบบดั้งเดิม การดำเนินการด้านการให้ความสนใจ (attention) จะถูกทำด้วยลำดับความยาวเดียวกันเสมอ ซึ่งให้รูปทรงคงที่สำหรับการทำงาน อย่างไรก็ตาม ในสถาปัตยกรรมแบบ dual-stream ที่เราพัฒนาขึ้นเอง บางชั้นของการให้ความสนใจจะดำเนินการเฉพาะกับความยาวของลำดับเท่านั้น ในขณะที่บางชั้นจะดำเนินการกับทั้งความยาวของลำดับและเงื่อนไข

แม้จะมีความท้าทายเหล่านี้ เราได้เห็นผลลัพธ์ที่น่าทึ่งหลังจากนำ CUDA Graphs มาใช้ เราใช้ค่าเวลาต่อโทเค็นผลลัพธ์ (TPOT) เพื่อวัดเวลาในการสร้างแต่ละโทเค็นระหว่างการอนุมาน หลังจากนำ CUDA Graphs มาใช้ ค่า TPOT ของเราดีขึ้นจาก 60.5 มิลลิวินาทีเป็น 20.5 มิลลิวินาที ซึ่งเพิ่มขึ้น 2.9 เท่า เวลาในการสร้างโดยรวมลดลง 66% จาก 31 วินาทีเหลือ 10.5 วินาที

เพื่อปรับปรุงความหน่วงเวลาให้ดียิ่งขึ้น เราได้นำการแคชแบบ KV มาใช้ ซึ่งเป็นแนวทางปฏิบัติมาตรฐานในการอนุมาน LLM ที่ได้รับการพิสูจน์แล้วว่ามีประสิทธิภาพสูงในอุตสาหกรรม 

ในแบบจำลองที่ใช้ทรานส์ฟอร์เมอร์ เช่น Cube 3D การสร้างโทเค็นแต่ละครั้งจำเป็นต้องคำนวณเมทริกซ์คีย์ (K) และเมทริกซ์ค่า (V) โดยอิงจากโทเค็นทั้งหมดที่สร้างขึ้นมาก่อนหน้า เมื่อลำดับยาวขึ้น การคำนวณเมทริกซ์เหล่านี้ใหม่สำหรับทุกโทเค็นจะยิ่งมีประสิทธิภาพน้อยลง

การแคช KV แก้ไขปัญหานี้โดย:

1. การเก็บเมทริกซ์ K และ V สำหรับโทเค็นทั้งหมดที่เคยสร้างขึ้นมาก่อนหน้านี้
2. การคำนวณเมทริกซ์ K และ V เฉพาะสำหรับโทเค็นใหม่เท่านั้น
3. การเพิ่มเมทริกซ์ใหม่เหล่านี้ไปยังค่าที่เก็บไว้ในแคช

แนวทางนี้ช่วยลดการคำนวณที่ซ้ำซ้อน ทำให้ลดปริมาณงานที่ต้องทำสำหรับแต่ละโทเค็นใหม่ สิ่งนี้จะมีผลกระทบมากขึ้นเมื่อลำดับที่สร้างขึ้นมีความยาวมากขึ้น

แนวทางของเราในการผสานรวมการแคช KV กับการใช้งาน CUDA Graph นั้นคล้ายคลึงกับการอนุมาน LLM แบบดั้งเดิม การเพิ่มการแคช KV ช่วยลด TPOT ของเราเหลือเพียง 7.8 มิลลิวินาที เวลาในการสร้างโดยรวมลดลง 87% จากเดิม 31 วินาที เหลือเพียง 4 วินาที การลดเวลาอย่างมีนัยสำคัญนี้ทำให้มีประสิทธิภาพมากขึ้นสำหรับผู้สร้างที่ใช้เครื่องมือนี้

เรากำลังสำรวจเทคนิคที่ช่วยลดความหน่วงและปรับปรุงประสบการณ์ของผู้ใช้ให้ดียิ่งขึ้น รวมถึงการใช้เคอร์เนลที่ปรับแต่งให้เหมาะสม การทำควอนไทซ์โมเดลเพื่อการอนุมานที่รวดเร็วขึ้น การปรับแต่งเฉพาะฮาร์ดแวร์ และการสร้างโทเค็นแบบขนาน

งานนี้มีความซับซ้อนมากขึ้นเมื่อเราขยายไปสู่การสร้างและเข้าใจฉากเต็มรูปแบบ ซึ่งองค์ประกอบ 3 มิติจำนวนมากต้องทำงานร่วมกันในบริบทเดียวกันภายในเลย์เอาต์ เราต้องการให้วัตถุและโลก 3 มิติที่เราสร้างขึ้นสามารถทำงานได้อย่างสมบูรณ์ เช่น ประตูเปิดและปิดได้ ล้อหมุนได้ ฯลฯ เพื่อให้บรรลุเป้าหมายนี้ เราจำเป็นต้องมีการสร้างและปรับปรุงอย่างรวดเร็วเพื่อขยายไปสู่ฉากทั้งหมด วัตถุที่ทำงานได้อย่างสมบูรณ์ และอวตาร เรารู้สึกตื่นเต้นที่จะได้แบ่งปันการปรับปรุงเพิ่มเติมและฟังก์ชันใหม่ ๆ ในขณะที่เรากำลังขยายโมเดลพื้นฐาน Cube 3D ของเรา และเฝ้าดูโลกเสมือนจริงที่ชุมชนผู้สร้างของเราสร้างขึ้นด้วยโมเดลเหล่านี้