ลดเวลา refit สูงสุด 5.7 เท่า เอ็นวิเดีย(NVDA) เปิดตัว Reshard

เอ็นวิเดีย(NVIDIA, NVDA) เปิดตัว Reshard collective ซึ่งลดเวลาแฝงในการส่งน้ำหนักโมเดลจาก GPU สำหรับฝึกไปยัง GPU สำหรับ inference ในงานการเรียนรู้แบบเสริมกำลัง (RL) ได้สูงสุด 5.7 เท่า โดยการทดสอบภายใน NeMo RL พบว่าเวลา refit สั้นลง 2.7 ถึง 5.7 เท่าเมื่อเทียบกับวิธีเดิม

เอ็นวิเดียแนะนำ operation ใหม่นี้และ NCCL M2N Reshard ซึ่งเป็นตัวอย่างการติดตั้งใช้งานแบบโอเพนซอร์สในบล็อกเทคนิคที่เผยแพร่เมื่อวันที่ 6 พฤศจิกายน ในขั้นตอน post-training ของ RL จะใช้ trajectory ที่ GPU สำหรับ inference สร้างขึ้นเพื่ออัปเดตโมเดลบน GPU สำหรับฝึก

เมื่อแยกงานฝึกกับ inference ไว้บน GPU คนละชุด จะต้องส่งน้ำหนักที่อัปเดตแล้วจาก GPU สำหรับฝึกไปยัง GPU สำหรับ inference เป็นระยะ กระบวนการนี้เรียกว่า weight refit และหากใช้เวลานานก็อาจทำให้รอบการอัปเดตโมเดลช้าลง

วิธี gather-and-broadcast แบบเดิมจะรวบรวมน้ำหนักที่แบ่งเก็บไว้บน GPU สำหรับฝึกให้เป็น tensor เต็มชุด แล้วกระจายไปยัง GPU สำหรับ inference ซึ่งจะเก็บเฉพาะส่วนข้อมูลที่ต้องใช้

Reshard คำนวณส่วนที่ทับซ้อนกันระหว่างการแบ่งข้อมูลต้นทางกับปลายทาง แล้วส่งเฉพาะข้อมูลที่จำเป็นโดยตรง วิธีนี้ลดขั้นตอนสร้าง tensor เต็มชุดและส่งข้อมูลซ้ำ

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

ในการทดสอบ NeMo RL บน DGX H100 รุ่น Qwen 3 ขนาด 235B ใช้เวลา refit ลดจาก 6.4 วินาทีเหลือ 1.9 วินาที ส่วน DeepSeek V3 ขนาด 671B ลดจาก 14.9 วินาทีเหลือ 2.6 วินาทีบนเครื่องเดียวกัน

บน DGX GB200 เวลา refit ของ DeepSeek V3 ขนาด 671B ลดจาก 6.1 วินาทีเหลือ 2.3 วินาที ขณะที่ Nemotron Ultra V3 ลดจาก 8.2 วินาทีเหลือ 2.3 วินาที

การเปรียบเทียบใช้วิธี gather-and-broadcast แบบเดิมที่เอ็นวิเดียเป็นผู้วัด ผลที่รายงานจำกัดอยู่ที่เวลาแฝงของ weight refit ไม่ได้หมายความว่าเวลาเรียนรู้ RL โดยรวม หรือ throughput จะดีขึ้นในสัดส่วนเดียวกัน

การทดสอบ microbenchmark แยกสำหรับ refit วัด Qwen3.8 ขนาด 2.4T ในรูปแบบ BF16 ได้ 4.3 วินาที ตัวเลขนี้วัดเฉพาะงาน refit

เอ็นวิเดียระบุว่าใน large language model รุ่นล่าสุด น้ำหนักของ feed-forward projection มีสัดส่วนประมาณ 95% ของพารามิเตอร์ทั้งหมด NeMo RL ใช้ Reshard กับส่วนนี้ ส่วนรายการที่เหลือ เช่น QKV projection และน้ำหนัก Mamba ยังคงใช้เส้นทางการทำงานเดิม

เมื่อวันที่ 22 กรกฎาคม NeMo RL เคยหารือเรื่องการปรับปรุง weight refit ระหว่าง GPU ที่ไม่ได้ใช้หน่วยความจำร่วมกัน การประกาศครั้งนี้เพิ่มเติมจากการหารือเดิมด้วยการเผยแพร่ตัวอย่างการติดตั้งใช้งาน Reshard และผลการวัดประสิทธิภาพ

NCCL M2N Reshard ถูกเผยแพร่เป็นไลบรารีทดลองในคลังโค้ดสาธารณะ คู่มือแพ็กเกจ Python ระบุว่าต้องใช้ NCCL 2.30.7 และยังไม่ได้ตรวจสอบกับ NCCL เวอร์ชันอื่น

เอกสาร vLLM มีตัวอย่าง M2N สำหรับย้าย weights ระหว่างการฝึกกับ inference โดยใช้ GPU H100 จำนวน 4 ตัว ตัวอย่างนี้แสดงวิธีใช้งาน แต่ไม่ได้เป็นการทดสอบซ้ำผลด้านประสิทธิภาพของเอ็นวิเดียโดยอิสระ

บล็อกเทคนิคของเอ็นวิเดียที่เผยแพร่เมื่อวันที่ 6 พฤศจิกายนระบุว่า NeMo RL กู้คืนจากความขัดข้องด้วยการทิ้ง communication group แล้วสร้างขึ้นใหม่ โดยนำเสนอขั้นตอนนี้เป็นส่วนหนึ่งของการติดตั้งใช้งานปัจจุบัน

ผู้ใช้รายหนึ่งรายงานใน issue สาธารณะ #23 ว่า เมื่อโปรเซสฝั่งตรงข้ามหยุดทำงาน คำสั่ง comm.abort() ไม่คืนค่าภายใต้ NCCL 2.30.7 และ nccl-extensions 0.1.0 ผู้รายงานระบุเพิ่มเติมว่าพบปัญหาซ้ำบน NCCL 2.32-3