การสื่อสารแบบ Mixture-of-Experts ที่มีประสิทธิภาพและพกพาได้

เรานำเสนอไลบรารีโอเพนซอร์สประสิทธิภาพสูงและพกพาได้สำหรับการสื่อสาร Mixture-of-Experts (MoE) ซึ่งให้ประสิทธิภาพเร็วกว่า 10 เท่าเมื่อเทียบกับพริมทิฟการสื่อสารแบบ All-to-All มาตรฐาน

การใช้งานของเรามีนวัตกรรมทางเทคนิคที่สำคัญหลายประการที่มอบประสิทธิภาพการสื่อสาร MoE ที่เหนือกว่า:

การสื่อสารที่เริ่มต้นโดย GPU (IBGDA): รองรับการสื่อสารโดยตรงจาก GPU ไปยัง NIC ซึ่งช่วยลดความหน่วงได้อย่างมากโดยไม่ต้องอาศัย CPU

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

ประสิทธิภาพในโหนดเดียวที่เร็วที่สุด: ความหน่วงต่ำกว่าการใช้งานที่เร็วที่สุดก่อนหน้านี้ 2.5 เท่าในการกำหนดค่าโหนดเดียว

ประสิทธิภาพข้ามโหนดที่มีประสิทธิภาพและพกพาได้: การใช้งานของเราบรรลุความเร็วสูงกว่าการสื่อสารแบบ all-to-all มาตรฐานสูงสุด 10 เท่า แม้ว่าจะช้ากว่าการใช้งานเฉพาะทางขั้นสูงประมาณ 2 เท่า แต่วิธีการของเราก็นำเสนอความสามารถในการพกพาที่ดีกว่าข้ามเวอร์ชัน NVSHMEM และสภาพแวดล้อมเครือข่าย (NVLink, CX-7 และ EFA)

ไลบรารีนี้เป็นโอเพนซอร์สเต็มรูปแบบและใช้งานได้ที่ https://github.com/ppl-ai/pplx-kernels

ในบทความนี้ เราจะสำรวจความท้าทายของความขนานของผู้เชี่ยวชาญในโมเดล MoE ขนาดใหญ่ และอธิบายแนวทางของเราในการส่งและรวมโทเค็นที่มีประสิทธิภาพทั่วทั้งสภาพแวดล้อม GPU แบบกระจาย

บทนำ

โมเดล Mixture-of-Experts (MoE) เช่น DeepSeek R1 และ Mixtral 8x7B ช่วยปรับปรุงโมเดลแบบหนาแน่นโดยจำกัดจำนวนน้ำหนักที่ถูกเปิดใช้งานสำหรับแต่ละโทเค็น แทนที่จะบรรลุกำลังพารามิเตอร์สูงด้วยการเพิ่มจำนวนชั้นหรือเพิ่มขนาดของชั้นเชิงเส้นใน Multi-Layer Perceptron (MLP) ของแต่ละชั้นตัวถอดรหัส โมเดล MoE จะแทนที่ MLP แบบดั้งเดิมด้วยผู้เชี่ยวชาญหลายคนและเราเตอร์ ตัวอย่างเช่น จากพารามิเตอร์ 671B ของ DeepSeek R1 มีเพียง 37B เท่านั้นที่ถูกคูณด้วยโทเค็นที่กำหนดในระหว่างการอนุมาน สิ่งนี้ช่วยลดปริมาณการคำนวณที่จำเป็นในการถอดรหัสโทเค็นได้อย่างมาก และปรับปรุงความหน่วงเมื่อเทียบกับโมเดลแบบหนาแน่น

โมเดล MoE มีความท้าทายเพิ่มเติมบางประการสำหรับการทำอินเฟอเรนซ์ เมื่อเทียบกับโมเดลแบบ dense แม้ว่า experts เองจะเป็นชั้น MLP ขนาดเล็ก แต่แต่ละ decoder layer จะมี router ที่ตัดสินใจว่า token จะถูกส่งไปยัง experts ใด โดยแต่ละ token จะถูกส่งไปยัง multiple experts โดยทั่วไป router จะเป็นชั้นเชิงเส้นขนาดเล็กที่สร้างการกระจายความน่าจะเป็น ปกติจะเลือก experts ที่มีคะแนนสูงสุดแบบ top-K และคำนวณ activation สุดท้ายเป็นค่าเฉลี่ยถ่วงน้ำหนัก โดยนำ output ของ experts มาบวกกันหลังคูณด้วยน้ำหนักที่ได้มาจากการกระจายความน่าจะเป็นของ router

เพื่อลดความหน่วงในระบบ MoE แบบกระจาย ความขนานสามารถนำมาใช้ประโยชน์ข้ามอุปกรณ์หลายเครื่องได้ แต่นี่ทำให้เกิดความท้าทายในการสื่อสาร โมเดลเช่น Mixtral 8x7B หรือ Llama-70B สามารถใส่ลงในอุปกรณ์ 8 เครื่องข้ามโหนดเดียวได้ โดยได้รับประโยชน์จากการเชื่อมต่อ NVLink ที่รวดเร็วและมีความหน่วงต่ำ (สูงสุด 900Gbps) อย่างไรก็ตาม โมเดลขนาดใหญ่จำเป็นต้องมีการแบ่งชิ้นส่วน (sharding) ผู้เชี่ยวชาญข้ามหลายโหนดโดยใช้ InfiniBand (สูงสุดที่ 400Gbps) ซึ่งทำให้เกิดความท้าทายด้านความหน่วงเพิ่มเติม

ขยาย

ในบทความนี้ เราสำรวจปัญหาของความขนานของผู้เชี่ยวชาญ โดยอธิบายการใช้งานเคอร์เนล GPU ของเราที่ส่งโทเค็นไปยังอุปกรณ์ที่มีผู้เชี่ยวชาญที่ได้รับมอบหมายอย่างมีประสิทธิภาพ และรวบรวมกลับคืนมาโดยการคำนวณผลรวมถ่วงน้ำหนักของเอาต์พุตของผู้เชี่ยวชาญ แม้ว่าจะมีโมเดลที่ปรับแต่งอย่างเข้มข้นซึ่งกำหนดเป้าหมายฮาร์ดแวร์เครือข่ายเฉพาะอยู่ แต่เราก็นำเสนอการใช้งานที่พกพาได้สูงซึ่งอาศัยชุดพริมทิฟ NVSHMEM ขั้นต่ำ แต่ยังคงให้ประสิทธิภาพที่เร็วกว่าการสื่อสารแบบ all-to-all มาตรฐานถึง 10 เท่า การใช้งานของเราบรรลุประสิทธิภาพระดับแนวหน้าในการกำหนดค่าโหนดเดียว ขณะเดียวกันก็รักษาความสามารถในการพกพาที่ยอดเยี่ยมในสภาพแวดล้อมเครือข่ายต่างๆ

การแบ่งชิ้นส่วนและความขนาน

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

ขยาย

ความขนานของผู้เชี่ยวชาญ (Expert Parallelism - EP) ดังแสดงในรูปแรก จะทำให้การคำนวณของผู้เชี่ยวชาญเท่านั้นที่เป็นแบบขนาน ผู้เชี่ยวชาญที่แตกต่างกันจะถูกกำหนดให้กับอุปกรณ์ที่แตกต่างกัน ซึ่งทำหน้าที่เก็บน้ำหนักของพวกเขา หลังจากการกำหนดเส้นทาง โทเค็นจะถูกส่งไปยังอุปกรณ์ที่เกี่ยวข้อง โดยผลลัพธ์จะถูกรวบรวมและสะสมหลังจากนั้น ความซับซ้อนของการกำหนดเส้นทางขึ้นอยู่กับระดับความขนานในส่วนอื่นๆ ของโมเดล: การจำลองชั้นอื่นๆ อาจกำจัดความจำเป็นในการกำหนดเส้นทางไปเลย เนื่องจากแต่ละเรดก์สามารถเลือกโทเค็นจากตารางการกำหนดเส้นทางที่จำลองไว้ในเครื่องได้ อย่างไรก็ตาม หากมีเพียงหนึ่งเรดก์ที่ทำการกำหนดเส้นทาง จำเป็นต้องมีการกระจาย (broadcast) เพื่อส่งโทเค็น, ดัชนี และน้ำหนักไปยังผู้เชี่ยวชาญของตน สุดท้าย การรวมแบบ all-gather หรือการกระจายแบบ all-to-all จะทำการซิงโครไนซ์โทเค็นเอาต์พุตกับเรดก์ใดก็ตามที่ดำเนินการโมเดลต่อไป การใช้งานดังกล่าวค่อนข้างง่ายเนื่องจาก torch มีพริมทิฟที่จำเป็นอยู่แล้ว แม้ว่าจะสามารถรับประโยชน์บางประการได้ด้วยการรวมการสะสมเข้ากับการดำเนินการรวบรวมที่รวบรวมโทเค็นจากผู้เชี่ยวชาญก็ตาม

ขยาย

ความขนานเฉพาะผู้เชี่ยวชาญไม่ได้ปรับขนาดตามอุดมคติ เนื่องจากโหนดในคลัสเตอร์อาจไม่ได้ใช้งานในขณะที่โมเดลกำลังรันชั้นที่ไม่ใช่ผู้เชี่ยวชาญ เช่น attention, norm และ sampling อย่างไรก็ตาม การคำนวณของชั้นเหล่านี้ ซึ่งส่วนใหญ่คือ attention ยังสามารถได้รับประโยชน์จากความขนานของเทนเซอร์ (Tensor Parallelism - TP) โมเดลส่วนใหญ่อาศัย attention แบบหลายหัว ซึ่งหมายความว่าหัว attention และการฉายภาพ Q, K และ V ที่สอดคล้องกันสามารถแบ่งชิ้นส่วนข้ามอุปกรณ์ได้ โดยการจำลองหรือรวบรวมชิ้นส่วนระหว่างชั้นต่างๆ หาก attention กระจายไปทั่วทุกอุปกรณ์ all-gather สามารถซิงโครไนซ์การกระตุ้นได้ ช่วยให้สามารถจำลองการกำหนดเส้นทางได้ ซึ่งต้องใช้พริมทิฟการซิงโครไนซ์ที่คล้ายคลึงกับกรณีความขนานของผู้เชี่ยวชาญเพื่อให้การใช้งานมีประสิทธิภาพ อย่างไรก็ตาม มีขีดจำกัดสำหรับความขนานในระดับนี้ เนื่องจากลดจำนวนหัว attention ต่ำกว่าเกณฑ์ที่กำหนดจะให้ผลตอบแทนที่ลดลง

ขยาย

เพื่อใช้ประโยชน์จากอุปกรณ์ทั้งหมดให้ดีที่สุดและรองรับความขนานของผู้เชี่ยวชาญในระดับที่สูงมาก สูงสุด 128 หรือ 256 GPU จึงจำเป็นต้องมีความขนานของข้อมูล (Data Parallelism - DP) ภายใต้รูปแบบนี้ อุปกรณ์ที่โมเดลถูกแบ่งออกจะถูกจัดกลุ่มเพื่อจัดการคำขอพร้อมกัน โดยคำนวณ attention และรักษา KV caches ที่ถูกแบ่งชิ้นส่วนข้ามกลุ่มในเครื่องของตน อินสแตนซ์หลายตัวของกลุ่มคู่ขนานเหล่านี้ทำงานร่วมกันในการประเมินผู้เชี่ยวชาญ โดยแต่ละกลุ่มโฮสต์ชุดย่อยที่แตกต่างกันของน้ำหนักผู้เชี่ยวชาญ ขึ้นอยู่กับจำนวนหัว attention กลุ่มหนึ่งมักจะครอบคลุมถึงขนาดของทั้งโหนด เนื่องจากการสื่อสารภายในโหนดเร็วกว่า ตัวอย่างเช่น ในรูปด้านบน เรดก์ DP หนึ่งตัวให้บริการคำขอสองรายการ A และ B อย่างอิสระ โดยจัดการ attention, norm และการเก็บบัญชีอื่น ๆ สำหรับคำขอ เรดก์ DP อื่นประมวลผลคำขอ C ที่แตกต่างกัน อย่างไรก็ตาม โหนดแรกโฮสต์ครึ่งหนึ่งของผู้เชี่ยวชาญ ในขณะที่อีกโหนดโฮสต์อีกครึ่งหนึ่ง ดังนั้นหลังจากกำหนดเส้นทาง โทเค็นจาก A และ B อาจถูกส่งไปยังโหนดที่สองและในทางกลับกัน สิ่งนี้นำไปสู่ปัญหาการสื่อสารแบบเบาบาง (sparse communication): แต่อุปกรณ์แต่ละเครื่องอาจส่งจำนวนโทเค็นที่แตกต่างกันไปยังปลายทางปลายทางใดก็ได้ พริมทิฟที่มีอยู่จาก torch โดยเฉพาะ all_to_all ไม่ค่อยเหมาะสมนัก เนื่องจากอาจต้องใช้รูปแบบของการแพดดิ้ง (padding) หรือการซิงโครไนซ์ GPU กับ CPU และการกระจายเมตาดาต้า เพื่อดำเนินการสื่อสารอย่างมีประสิทธิภาพ จำเป็นต้องมีเคอร์เนลที่กำหนดเองเพื่อประมวลผลข้อมูลการกำหนดเส้นทางและเริ่มต้นการสื่อสารจากอุปกรณ์โดยไม่ต้องอาศัย CPU หลังจากกำหนดเส้นทาง เคอร์เนลการจัดส่ง (dispatch kernel) จะต้องส่งโทเค็นไปยังเรดก์ที่พวกมันถูกกำหนดเส้นทางไป ในขณะที่ฝั่งรวม (combine side) จะต้องรวบรวมการกระตุ้นที่เป็นของคำขอในกลุ่ม DP ปัจจุบัน นอกจากนี้ งานจะต้องถูกปรับสมดุลภายในกลุ่ม DP ซึ่งอาจมีอุปกรณ์หลายเครื่อง เพื่อลดการกระจายให้น้อยที่สุดและจำลองโทเค็นบนเรดก์ทั้งหมดอย่างถูกต้องสำหรับชั้นถัดไป

NVSHMEM

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

การดำเนินการ NVSHMEM ถูกสร้างขึ้นรอบแนวคิดของหน่วยความจำสมมาตร (symmetric memory): พวกมันทำงานบนบัฟเฟอร์ที่ได้รับการจัดสรรบนอุปกรณ์ทั้งหมดที่เข้าร่วมในการสื่อสารระหว่างอุปกรณ์ อุปกรณ์ในเครื่องสามารถใช้ที่อยู่ซึ่งได้มาจากบัฟเฟอร์ในเครื่องของตนเองควบคู่ไปกับดัชนีของเรดก์เป้าหมายเพื่อระบุปลายทางของการดำเนินการ รูปด้านล่างแสดงแนวคิดนี้: GPU ทั้งสองจัดสรรบัฟเฟอร์สมมาตรที่มีขนาดเท่ากัน โดยเก็บพอยเตอร์ src ไว้ เรดก์แรกต้องการส่งจำนวนเต็ม 3 ตัวไปยังเรดก์ที่สอง โดยวางไว้ที่จุดเริ่มต้นของบัฟเฟอร์: nvshmem_int_put_nbi ดึงที่อยู่เริ่มต้นจากบัฟเฟอร์ในเครื่อง โดยระบุอุปกรณ์เป้าหมาย เรดก์ที่สองดึงออฟเซ็ตจากบัฟเฟอร์ของตนเอง โดยส่งหนึ่งองค์ประกอบไปยังอุปกรณ์แรก โดยเลื่อนออฟเซ็ตทีละหนึ่ง แม้ว่าที่อยู่ปลายทางจะต้องเป็นบัฟเฟอร์สมมาตรที่จัดสรรโดยใช้ nvshmem_alloc เสมอ แต่บัฟเฟอร์ต้นทางสามารถเป็นภูมิภาคใดก็ได้ของหน่วยความจำอุปกรณ์ โดยมีเงื่อนไขว่าได้ลงทะเบียนล่วงหน้ากับ NVSHMEM แล้ว

แม้ว่า NVSHMEM จะมีพริมทิฟที่หลากหลาย แต่เคอร์เนลของเราอาศัยฟังก์ชันเพียง 3 ฟังก์ชันเท่านั้น โดยสร้างการซิงโครไนซ์และการฟันซิง (fencing) ทั้งหมดขึ้นมาจากฟังก์ชันเหล่านี้

ขยาย

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

nvshmemx_signal_op ทำงานบนตำแหน่งหน่วยความจำเดี่ยว โดยทั่วไปคือแฟล็ก 64 บิต โดยตั้งค่าหรือเพิ่มค่าแบบอะตอมมิค มีประโยชน์ในการส่งเมตาดาต้าและซิงโครไนซ์อุปกรณ์

nvshmem_uint64_wait_until ใช้ในฝั่งรับของสัญญาณเพื่อสำรวจ (poll) แฟล็กจนกว่าอุปกรณ์ระยะไกลจะอัปเดต

การรักษาความถูกต้องของลำดับใดๆ ระหว่างการดำเนินการผ่าน nvshmem_quiet หรือ nvshmem_fence นั้นมีค่าใช้จ่ายสูง ดังนั้นการใช้งานเคอร์เนลของเราจึงหลีกเลี่ยงการใช้พริมทิฟ NVSHMEM อื่นๆ และลดการพึ่งพานอกเหนือจากคู่ส่ง-รับแต่ละคู่ เมื่อส่งข้อมูล เรามักจะใช้ฟังก์ชันเวอร์ชันที่ไม่บล็อก (non-blocking) เสมอ โดยไม่ต้องรอให้ข้อมูลถูกส่งออกจากเรดก์ในเครื่องด้วยซ้ำ การซิงโครไนซ์อื่นๆ โดยปริยายข้ามกลุ่มเคอร์เนลช่วยให้แน่ใจว่าบัฟเฟอร์เหล่านี้จะไม่ถูกเขียนทับในลักษณะทำลายล้างในขณะที่ NIC ต้องการเข้าถึง

NVSHMEM ใช้ GPUDirect RDMA เพื่อให้ NIC สามารถอ่านและเขียนหน่วยความจำ GPU ได้โดยตรงสำหรับการแลกเปลี่ยนข้อมูลโดยไม่ต้องเกี่ยวข้องกับ CPU นอกจากนี้ บน NIC ของ ConnectX แล้ว NVSHEMEM ยังรองรับการสื่อสารที่เริ่มต้นโดย GPU (หรือที่เรียกว่า Infiniband GPUDirect Async หรือ IBGDA) ซึ่งช่วยให้ GPU สามารถโพสต์การดำเนินการ RDMA ไปยัง NIC ได้โดยตรงและสำรวจคิวความสำเร็จของ NIC ได้โดยตรง บนแพลตฟอร์มที่ไม่รองรับการสื่อสารที่เริ่มต้นโดย GPU NVSHMEM จะสร้างเธรด “พร็อกซี” บน CPU ของโฮสต์เพื่อเริ่มต้นการสื่อสารและสำรวจความสำเร็จในนามของ GPU โปรแกรม NVSHMEM สามารถพกพาได้ไม่ว่าจะรองรับการสื่อสารที่เริ่มต้นโดย GPU หรือไม่ก็ตาม อย่างไรก็ตาม การสื่อสารที่เริ่มต้นโดย GPU จะช่วยลดความหน่วงลงอย่างมากเนื่องจากข้ามการอ้อมไปยัง CPU ไปโดยสิ้นเชิง

เคอร์เนลแบบพกพา

เราใช้งานการสื่อสาร MoE ผ่านคู่ของเคอร์เนลส่ง (dispatch) และรวม (combine) เคอร์เนลส่งมีหน้าที่อ่านโทเค็นและข้อมูลการกำหนดเส้นทางบนแต่ละเรดก์ โดยส่งไปยังผู้เชี่ยวชาญที่เหมาะสม เคอร์เนลรวมจะรวบรวมการกระตุ้นที่ผลิตโดยผู้เชี่ยวชาญและส่งพวกมันกลับไปยังเรดก์ต้นทาง พร้อมทั้งคำนวณค่าเฉลี่ยถ่วงน้ำหนักจากผู้เชี่ยวชาญที่เลือกตามน้ำหนักที่คำนวณโดยเราเตอร์ เคอร์เนลยังถูกแบ่งออกเป็นส่วนส่งและรับ เพื่อให้สามารถทับซ้อนการถ่ายโอนข้อมูลกับการคำนวณได้ เคอร์เนลส่งทั้งหมดไม่บล็อกและไม่ซิงโครไนซ์: พวกมันเพียงแค่ส่งการเขียนทั้งหมดไปยังระยะไกล ในทางกลับกัน เคอร์เนลรวมจะอ่านจากหน่วยความจำเท่านั้น โดยรอจนกว่าข้อมูลที่จำเป็นทั้งหมดจะถูกถ่ายโอน หลังจากส่งงานไปยัง NIC ในขณะที่ข้อมูลถูกถ่ายโอนแบบอะซิงโครนัสผ่านสายเคเบิล GPU สามารถทำงานที่มีประโยชน์อื่นๆ ในเครื่องได้ เช่น การใช้ผู้เชี่ยวชาญที่แชร์หรือการคำนวณ attention

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

ขยาย

รูปแบบเดียวของการซิงโครไนซ์ทั่วโลก ดังแสดงในรูปด้านบน ถูกนำมาใช้ในเคอร์เนลรวม-รับ (combine-receive) เมื่อข้อมูลถูกส่งจากฝั่งส่ง เรดก์จะส่งคืนและดำเนินการต่อเมื่อได้รับโทเค็นทั้งหมดที่ต้องการจากเรดก์อื่นๆ และเข้าสู่เคอร์เนลรวม-ส่งทันทีที่ได้รับการประมวลผล แบร์ริเยร์ (barrier) ในเคอร์เนลรวม-รับช่วยให้แน่ใจว่าไม่มีเรดก์ใดสามารถวิ่งนำหน้าและเริ่มส่ง-ส่ง (dispatch-send) ในขณะที่เรดก์อื่นยังคงรอรับข้อมูล ซึ่งอาจทำให้เกิดการทับซ้อนที่เป็นอันตราย การซิงโครไนซ์ทำได้ในเคอร์เนลรวมเพื่อความเรียบง่าย โดยเป็นการเพิ่มประสิทธิภาพ แบร์ริเยร์เดียวกันนี้อาจถูกนำมาใช้ระหว่างเคอร์เนลรับ-ส่ง (dispatch receive) และเคอร์เนลรวม-ส่ง (combine send) ซึ่งช่วยลดความหน่วงลงไปอีกโดยให้เวลากับ NIC มากขึ้นในการสะ้างการถ่ายโอนในขณะที่ GPU กำลังรันเคอร์เนลอื่น

เคอร์เนลทั้งสองถูกแบ่งออกเป็น SM ทั้งหมดที่มีอยู่ของอุปกรณ์: ในขณะที่เคอร์เนลส่ง-ส่ง (dispatch send) และรวม-รับ (combine receive) ต้องทำแบบขนานตามจำนวนโทเค็นสูงสุดต่อเรดก์ (max_m) เคอร์เนลรับ-ส่ง (dispatch receive) และรวม-ส่ง (combine send) ต้องจัดการโทเค็นจำนวนมากจากแต่ละเรดก์และผู้เชี่ยวชาญในท้องถิ่น โดยมีค่าสูงสุดคือ max_m * (num_experts // EP) * (EP // TP)

การส่ง (Dispatch)

เคอร์เนลการส่งมีหน้าที่ส่งโทเค็นและปัจจัยการปรับขนาดบล็อกไปยังเพียร์ที่เป็นโฮสต์ของผู้เชี่ยวชาญที่โทเค็นถูกส่งไป ฝั่งผู้ส่งอาศัยความเชี่ยวชาญของวาร์ป (warp specialization) เพื่อทำภารกิจสองอย่างให้เป็นคู่ขนาน: การรวบรวมข้อมูลการกำหนดเส้นทางเพื่อพิจารณาว่ามีกี่โทเค็นที่ถูกส่งจากแต่ละเรดก์ และการบรรจุโทเค็นลงในข้อความที่จะส่งไปยังระยะไกล ฝั่งผู้รับจะรอให้จำนวนโทเค็นทั้งหมดถูกรับมาก่อน จากนั้นจึงอ่านและแกะเพย์โหลดจากหน่วยความจำร่วมลงในเทนเซอร์ที่จะถูกส่งต่อไปยังเคอร์เนล Grouped GEMM ที่อยู่ปลายน้ำ นอกจากนี้ยังจัดเก็บข้อมูลลงในบัฟเฟอร์ที่แชร์กับเคอร์เนลรวม เพื่อระบุว่าควรส่งโทเค็นแต่ละรายการกลับไปที่ใด ข้อมูลนี้จำเป็นเนื่องจากเคอร์เนลรับสลับตำแหน่งโทเค็นในบัฟเฟอร์ที่ต่อเนื่องกันในลำดับที่ไม่ใช่แบบกำหนดรูป (non-deterministic)

ขยาย

ในส่วนของผู้ส่ง วาร์ปเดี่ยวที่มี 32 เธรดในแต่ละบล็อกมีหน้าที่อ่านข้อมูลการกำหนดเส้นทางและนับจำนวนโทเค็นที่ผู้เชี่ยวชาญปลายทางแต่ละรายคาดว่าจะได้รับจากเรดก์นี้ โดยส่งจำนวนบวกหนึ่งโดยใช้ nvshmem_signal_op จำนวนจะถูกเพิ่มขึ้นทีละหนึ่ง เนื่องจากเปลี่ยนจากศูนย์เป็นไม่ใช่ศูนย์ที่ฝั่งปลายทางหมายถึงการได้รับจำนวนจากเรดก์ต้นทาง ในเวลาเดียวกัน วาร์ปที่เหลือจะร่วมมือกันเพื่อคัดลอกโทเค็นไปยังหน่วยความจำสมมาตรข้ามบล็อกทั้งหมดแบบขนาน โดยบรรจุการกระตุ้น, ปัจจัยการปรับขนาด และดัชนีของพวกมันบนเรดก์ในเครื่องลงในชิ้นส่วนต่อเนื่อง ดัชนีจำเป็นสำหรับผู้ส่งรวมเพื่อกำหนดที่อยู่ที่จะเขียนโทเค็นลงไป ถัดไป หลังจากตรวจสอบให้แน่ใจว่าข้อมูลทั้งหมดถูกคัดลอกแล้วผ่านแบร์ริเยร์ข้ามกลุ่มวาร์ป วาร์ปจะยังคงทำงานอย่างอิสระอีกครั้ง โดยแต่ละวาร์ปจะส่งบัฟเฟอร์เดียวกันไปยังผู้เชี่ยวชาญที่แตกต่างกันแบบขนาน โทเค็นจะถูกส่งโดยใช้ nvshmemx_putmem_signal_nbi_warp ซึ่งจะเพิ่มจำนวนโทเค็นที่ส่งจากเรดก์ในเครื่องบนอุปกรณ์ระยะไกลแบบอัตโนมัติ ภายในกลุ่ม DP เนื่องจากเรดก์แต่ละตัวเป็นเจ้าของสำเนาของโทเค็นที่จะส่ง การส่งจึงถูกปรับสมดุลอย่างเท่าเทียมกัน โดยอุปกรณ์แต่ละเครื่องจะส่งชุดย่อยของโทเค็นออกไป

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

การรวม (Combine)

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

ขยาย

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

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

การประเมินประสิทธิภาพ (Benchmarks)

เราประเมินเคอร์เนลของเราบนคลัสเตอร์ของ 128 GPU ที่กระจายข้าม 16 โหนดที่เชื่อมต่อผ่าน InfiniBand และ CX-7 NIC เราลองทั้ง GPUDirect Async (IBGDA) และ Reliable Connection (RC) กับพร็อกซี CPU เราเปรียบเทียบกับสถาปัตยกรรม DeepSeek ที่ปรับแต่งอย่างสูง รวมถึงพริมทิฟแบบหนาแน่นที่จัดเตรียมโดย PyTorch (ผ่าน NCCL) หรือ NVSHMEM

ขยาย

แม้ว่าในการดำเนินการ all-to-all ของ NVSHMEM แบบหนาแน่น ประสิทธิภาพของ IBRC และ IBGDA จะคล้ายคลึงกัน (6378 µs เทียบกับ 6180 µs) แต่ IBGDA จะเร็วกว่าอย่างเห็นได้ชัดกับเคอร์เนลแบบเบาบาง เมื่อรวม Dispatch และ Combine เข้าด้วยกัน IBGDA ใช้เวลา 902 µs ในขณะที่ IBRC ใช้เวลา 3223 µs ซึ่งเป็นการปรับปรุงความหน่วงถึง 3.6 เท่า ในขณะที่ all-to-all ถูกจำกัดด้วยแบนด์วิดท์ แต่เคอร์เนลแบบเบาบางจะกระจายข้อมูลน้อยกว่าหลายสัปดาห์ โดยถูกจำกัดด้วยความหน่วง ด้วยการทริกเกอร์การถ่ายโอนเครือข่ายโดยตรงจาก GPU โดยไม่ต้องอาศัยพร็อกซี CPU มาประสานงานระหว่าง GPU และ NIC ความหน่วงแบบ end-to-end จึงลดลงอย่างมาก

ในขณะที่เคอร์เนลแบบพกพาของเราช้ากว่าเคอร์เนล DeepSeek ที่ปรับแต่งอย่างสูงประมาณ 2 เท่า แต่พวกมันช่วยปรับปรุงความหน่วงได้ 10 เท่าเมื่อเทียบกับเคอร์เนลแบบหนาแน่น (902 µs เทียบกับ 9944 µs)

นอกจากนี้ การแบ่งออกเป็นส่วนผู้ส่ง-ผู้รับยังช่วยให้สามารถซ่อนความหน่วงบางส่วนได้ ซึ่งแตกต่างจากพริมทิฟของไลบรารี

ขยาย

บนโหนดเดียว (EP8) NVSHMEM ใช้ NVLINK สำหรับการขนส่ง ซึ่งให้ความหน่วงต่ำกว่าและปริมาณงานสูงกว่าเครือข่ายระหว่างโหนด เคอร์เนลแบบพกพาของเราเร็วกว่า DeepEP บนโหนดเดียวประมาณ 2.5 เท่า (186 µs เทียบกับ 481 µs)

การปรับปรุงเพิ่มเติม

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

บทสรุป

เราได้นำเสนอไลบรารีประสิทธิภาพสูงและพกพาได้สำหรับการสื่อสาร MoE ซึ่งบรรลุประสิทธิภาพเร็วกว่า 10 เท่าเมื่อเทียบกับการสื่อสารแบบ all-to-all มาตรฐาน พร้อมทั้งรักษาความเข้ากันได้ข้ามการกำหนดค่าฮาร์ดแวร์ที่หลากหลาย บนการปรับใช้โหนดเดียวด้วย NVLink โซลูชันของเราแสดงให้เห็นความหน่วงที่ต่ำกว่าการใช้งานก่อนหน้านี้ 2.5 เท่า

แนวทางของเราสร้างความสมดุลระหว่างประสิทธิภาพและความสามารถในการพกพาผ่านนวัตกรรมที่สำคัญ ซึ่งรวมถึงการสนับสนุนการสื่อสารที่เริ่มต้นโดย GPU, สถาปัตยกรรมเคอร์เนลแยกที่ช่วยให้สามารถทับซ้อนการคำนวณและการสื่อสารได้ และการส่งโทเค็นที่มีประสิทธิภาพโดยใช้พริมทิฟ NVSHMEM ขั้นต่ำ แม้ว่าจะช้ากว่าการใช้งานเฉพาะทางขั้นสูงบนการตั้งค่าหลายโหนดประมาณ 2 เท่า แต่ไลบรารีของเราก็มีความยืดหยุ่นที่เหนือกว่าในสภาพแวดล้อมเครือข่ายต่างๆ (NVLink, CX-7 และ EFA)

เนื่องจากโมเดล MoE ยังคงปรับขนาดต่อไป กลยุทธ์การสื่อสารที่มีประสิทธิภาพเช่นเราจะมีความสำคัญมากขึ้นเรื่อยๆ สำหรับการปรับใช้งานจริง การใช้งานแบบโอเพนซอร์สเต็มรูปแบบของเรามีอยู่ที่ https://github.com/ppl-ai/pplx-kernels

เอกสารอ้างอิง

DeepEP

รายงานทางเทคนิคของ DeepSeek-V3

NVIDIA NVSHMEM

การอ้างอิง

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}