การถ่ายโอนน้ำหนักสำหรับการฝึกโพสต์ RL ในเวลาน้อยกว่า 2 วินาที

เราเพิ่งบรรลุการอัปเดตพารามิเตอร์ข้ามเครื่องในเวลาเพียง 1.3 วินาที สำหรับ Kimi-K2 (1 ล้านล้านพารามิเตอร์) โดยถ่ายโอนน้ำหนักจาก GPU สำหรับฝึก 256 ตัว (BF16) ไปยัง GPU สำหรับอนุมาน 128 ตัว (FP8)

ในการปรับแต่งการเรียนรู้แบบเสริมแรงแบบอะซิงโครนัส (asynchronous reinforcement learning fine-tuning) การฝึกและการอนุมานจะทำงานบน GPU แยกกัน หลังจากแต่ละขั้นตอนการฝึก จะต้องพุชน้ำหนักใหม่ไปยังโหนดอนุมาน เฟรมเวิร์กที่มีอยู่จำนวนมากใช้เวลาหลายวินาที หรือแม้แต่หลายนาที สำหรับโมเดลระดับล้านล้านพารามิเตอร์

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

โซลูชันของเราสร้างขึ้นบน RDMA WRITE ซึ่งเป็นฟังก์ชันดั้งเดิมแบบทางเดียว (one-sided primitive) โดยที่ต้นทางจะเขียนข้อมูลลงในหน่วยความจำ GPU ปลายทางโดยตรง

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

ฝ่ายปลายทางจะไม่ได้รับการแจ้งเตือนสำหรับการถ่ายโอนนี้ด้วยซ้ำ สิ่งนี้ทำให้เราได้รับการถ่ายโอนที่มีความหน่วงต่ำ (low-latency) มีปริมาณงานสูง (high-throughput) และไม่ต้องคัดลอกซ้ำ (zero-copy) ซึ่งขับเคลื่อนโดยโหนดฝึกโดยไม่ต้องมีลอจิกควบคุมใดๆ บนโหนดอนุมาน

การรวบรวมข้อมูลเมตา – ตัวควบคุมจะรวบรวมข้อมูลเมตาของพารามิเตอร์จาก GPU สำหรับฝึกและอนุมานทั้งหมด

การคำนวณกำหนดการ – ตัวควบคุมจะคำนวณกำหนดการถ่ายโอนน้ำหนักแบบคงที่ โดยแมปว่า GPU สำหรับฝึกตัวใดส่งพารามิเตอร์ใดไปยัง GPU สำหรับอนุมานตัวใด และตามลำดับใด

การกระจายกำหนดการ – ตัวควบคุมจะส่งกำหนดการไปยัง GPU สำหรับฝึกทั้งหมด

การดำเนินการ – หลังจากแต่ละขั้นตอนการฝึก ตัวควบคุมจะส่งสัญญาณให้ GPU สำหรับฝึกเริ่มการถ่ายโอน

เมื่อกำหนดเวิร์กโฟลว์ระดับสูงแล้ว ความท้าทายหลักคือวิธีดำเนินการถ่ายโอนน้ำหนักอย่างมีประสิทธิภาพในสเกลระดับล้านล้านพารามิเตอร์ ในที่นี้เราจะอธิบายรายละเอียดของเส้นทางการดำเนินการ

พารามิเตอร์ในการฝึกจะถูกกระจายไปตามตำแหน่ง FSDP โดยการใช้ full_tensor() ทำให้ GPU ทั้งหมดใน DeviceMesh สามารถสร้างพารามิเตอร์เต็มรูปแบบขึ้นใหม่ได้ ดังนั้นทั้งหมดจึงสามารถทำหน้าที่เป็นแหล่งที่มาสำหรับการถ่ายโอนน้ำหนักได้

DeviceMesh หลายตัวที่แยกออกจากกันจะสร้าง mesh group เนื่องจาก DeviceMesh ในกลุ่มเดียวกันแยกจากกัน การถ่ายโอนของพวกมันจึงไม่รบกวนกันและสามารถทำงานแบบขนานได้อย่างเต็มที่ ระหว่าง mesh group เราแทรก Global Barrier เพื่อบังคับใช้ลำดับ

ไปป์ไลน์การถ่ายโอนน้ำหนักทำให้การคัดลอกจากโฮสต์ไปยังอุปกรณ์ การประกอบเทนเซอร์ การหลอมรวมและการควอนไทซ์ และการถ่ายโอน RDMA ข้าม MeshGroups ซ้อนทับกัน

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

Host to device memcpy — หาก FSDP offload น้ำหนักไปยัง CPU

การเตรียมพารามิเตอร์ — สร้างน้ำหนักเต็มรูปแบบด้วย full_tensor() ใช้ projection fusion และทำการ quantize หากจำเป็น

การถ่ายโอนผ่าน RDMA — เขียนข้อมูลแบบ zero-copy ไปยังหน่วยความจำ GPU ระยะไกลสำหรับการอนุมาน

Global barrier — หลังจากเรียกใช้ full_tensor() ทั้งหมดเสร็จสิ้น ให้ซิงโครไนซ์ข้าม mesh group โดยใช้ GLOO ผ่านอีเทอร์เน็ต

ในการใช้งานจริง เราจะรักษาคิว FIFO ของงานสำหรับแต่ละขั้นตอนของไปป์ไลน์ เมื่อใดก็ตามที่งานที่อยู่หัวคิวทำขั้นตอนนั้นเสร็จสิ้น งานจะถูกย้ายไปยังท้ายคิวของขั้นตอนถัดไป

full_tensor() และการดำเนินการอื่นๆ ของ GPU ทำให้เกิดการใช้หน่วยความจำ GPU เพิ่มเติม เพื่อหลีกเลี่ยงข้อผิดพลาด out of memory เราจะเริ่มดำเนินการงานเฉพาะเมื่อเนื้องานที่กำลังดำเนินการอยู่ในปัจจุบัน (on-the-fly tasks) ใช้หน่วยความจำ GPU ชั่วคราวน้อยกว่าค่า watermark ที่กำหนดค่าไว้

ทางเลือกในการออกแบบหลายประการทำให้ระบบของเราทำงานได้เร็วกว่าอย่างเห็นได้ชัดและบำรุงรักษาง่ายกว่าโซลูชันโอเพนซอร์สทั่วไป

การถ่ายโอนน้ำหนักที่อิง Rank0 จะรวบรวมน้ำหนักก่อนแจกจ่ายไปยังโหนดอินเฟอเรนซ์

การสื่อสารแบบจุดต่อจุดจะถ่ายโอนน้ำหนักจากโหนดเทรนนิงทั้งสองไปยังอินเทอร์เฟซเครือข่ายของโหนดอินเฟอเรนซ์

รูปแบบทั่วไปคือการส่งพารามิเตอร์ทั้งหมดผ่าน GPU แบงก์ 0 (rank-0): รวมข้อมูลที่แบงก์ 0 ของการฝึก, ส่งไปยังแบงก์ 0 ของการอนุมาน, แล้วกระจายออกอีกครั้ง สิ่งนี้จะกลายเป็นจุดคอขวดอย่างรวดเร็ว ซึ่งจำกัดด้วยแบนด์วิดท์ PCIe ของ GPU ตัวเดียวและ NIC (เช่น 400 Gbps ≈ 50 GB/s)

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

การถ่ายโอนแบบ RPC ใช้บัฟเฟอร์ฝั่งอินเฟอเรนซ์และการเรียกของเวิร์กเกอร์ ขณะที่การถ่ายโอนแบบ one-sided จะเขียนน้ำหนักลงโดยตรงโดยไม่มีการมีส่วนร่วมของเวิร์กเกอร์ฝั่งอินเฟอเรนซ์

ระบบบางระบบอาศัยการเรียกใช้เมธอด update_weight() ของเอนจินการอนุมานสำหรับเทนเซอร์แต่ละตัว นั่นหมายถึงการเปลี่ยนแปลงที่เจาะลึกในโค้ดการอนุมาน บวกกับโอเวอร์เฮดจาก RPC, การซีเรียลไลซ์ (serialization), และการประสานงานของระนาบควบคุม (control-plane coordination)

ด้วยฟังก์ชันดั้งเดิม RDMA WRITE เราอัปเดตน้ำหนักบนหน่วยความจำ GPU ของการอนุมานได้อย่างเงียบๆ โดยไม่มีการคัดลอกเพิ่มเติม ไม่มีข้อความระนาบควบคุมและไม่ต้องใช้ลอจิกควบคุมของ CPU ไม่จำเป็นต้องมีการปรับเปลี่ยนเอนจินการอนุมานใดๆ

กระบวนการถ่ายโอนน้ำหนักสามารถใช้ประโยชน์จากทรัพยากรฮาร์ดแวร์สี่ประเภท: (1) การเคลื่อนย้ายข้อมูลระหว่างโฮสต์และอุปกรณ์ (2) การคำนวณของ GPU สำหรับ projection fusion และ quantization (3) เครือข่าย RDMA สำหรับระนาบข้อมูล (data plane) (4) อีเทอร์เน็ตสำหรับระนาบควบคุม (control plane)

การออกแบบของเราแบ่งงานถ่ายโอนน้ำหนักออกเป็นขั้นตอนไปป์ไลน์ ช่วยให้สามารถเหลื่อมการทำงานข้ามทรัพยากรฮาร์ดแวร์ต่างๆ ได้อย่างง่ายดาย

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

กำหนดการของเราถูกคำนวณเพียงครั้งเดียวตอนเริ่มต้นระบบ การวนซ้ำการฝึกแต่ละครั้งจะทำซ้ำตามแผน: ตัวควบคุมจะออกสัญญาณ "ไป" และ GPU จะปฏิบัติตามเส้นทางที่กำหนดไว้ล่วงหน้า การดำเนินการมีความแน่นอนและใช้ทรัพยากรน้อย

เป็นเรื่องง่ายที่จะดึงกระบวนการอัปเดตน้ำหนักทั้งหมดมารวมกันในฟังก์ชันเดียวแบบเสาหิน (monolithic function): รวบรวมข้อมูลเมตา, การจับคู่ชื่อ, การรวมข้อมูลภายในโหนด, projection fusion, quantization, การแบ่งย่อยเครือข่ายการสื่อสาร, และการถ่ายโอนเครือข่ายข้ามโหนด ซึ่งยากต่อการเขียนโปรแกรมให้ถูกต้อง และยากยิ่งกว่าในการเพิ่มประสิทธิภาพ

ในทางวิศวกรรมของเรา เราแยกขั้นตอนเหล่านี้ออกเป็นคอมโพเนนต์อิสระ คอมโพเนนต์แต่ละตัวสามารถทดสอบยูนิตเทสต์ (unit test) ทำความเข้าใจ และปรับประสิทธิภาพได้อย่างโดดเดี่ยว

การถ่ายโอนน้ำหนักที่รวดเร็วและเชื่อถือได้คือองค์ประกอบสำคัญสำหรับการปรับแต่ง RL สเกลใหญ่ ด้วยการรวม ฟังก์ชันดั้งเดิม RDMA WRITE, กำหนดการถ่ายโอนแบบคงที่, และ การดำเนินการแบบไปป์ไลน์ เราจึงลดเวลาการอัปเดตระดับล้านล้านพารามิเตอร์เหลือเพียง 1.3 วินาที บน Kimi-K2 แนวทางนี้เข้าใจง่าย บำรุงรักษาสะดวก และหลีกเลี่ยงคอขวดของการออกแบบแบบดั้งเดิม

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: การถ่ายโอนแบบทางเดียว

เวิร์กโฟลว์ระดับสูง

การดำเนินการถ่ายโอนน้ำหนัก

DeviceMesh และ Mesh Groups

ท่อส่งงาน (Task pipeline)

การควบคุมการใช้หน่วยความจำ GPU

ทำไมจึงรวดเร็วและเรียบง่าย

การสื่อสารแบบจุดต่อจุด (Point-to-point communication)

การถ่ายโอนข้อมูลแบบทางเดียว (One-sided data transfer)

การทำไปป์ไลน์ (Pipelining)

กำหนดการแบบคงที่ (Static Schedule)

การแยกส่วนที่สะอาดตา (Clean separation)

บทสรุป

การอ้างอิง