การถ่ายโอนหน่วยความจำ GPU ประสิทธิภาพสูงบน AWS Sagemaker Hyperpod

โครงสร้างพื้นฐานด้านการเรียนรู้เชิงลึกสมัยใหม่มักกำหนดให้มีการถ่ายโอนข้อมูลจำนวนมากระหว่าง GPU ข้ามเครื่องจักร ที่ Perplexity เราพบความท้าทายทางเทคนิคที่ไม่เหมือนใคร นั่นคือการถ่ายโอนส่วนหน่วยความจำ GPU ที่ไม่ต่อเนื่องกันระหว่างเครื่องจักรได้อย่างมีประสิทธิภาพด้วยความเร็วสูงสุดเท่าที่จะเป็นไปได้ แพลตฟอร์มเป้าหมายของเราคืออินสแตนซ์ AWS p5 ซึ่งมอบแบนด์วิดท์เครือข่ายที่น่าประทับใจถึง 3200 Gbps ผ่านการ์ดเครือข่าย 32 ใบ บทความนี้แบ่งปันเส้นทางของเราในการสร้างโซลูชันเครือข่ายประสิทธิภาพสูงแบบกำหนดเองที่บรรลุแบนด์วิดท์ทางทฤษฎีนี้ถึง 97.1%

ความท้าทายทางเทคนิค

กรณีการใช้งานของเราแสดงให้เห็นข้อกำหนดทางเทคนิคที่สำคัญหลายประการ:

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

แม้ว่าไลบรารี NCCL ของ NVIDIA จะเป็นมาตรฐานโดยพฤตินัยสำหรับการเรียนรู้เชิงลึกแบบกระจาย แต่ก็ยังไม่เหมาะสำหรับกรณีการใช้งานของเรา:

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

เครือข่ายประสิทธิภาพสูงสมัยใหม่

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

เครือข่ายส่วนใหญ่ที่เราใช้ทุกวันอาศัยโปรโตคอล TCP/IP ซึ่งแอปพลิเคชันสื่อสารกับการ์ดเครือข่ายผ่านเคอร์เนลของระบบปฏิบัติการโดยใช้ซ็อกเก็ต อย่างไรก็ตาม เครือข่ายประสิทธิภาพสูงใช้ RDMA (Remote Direct Memory Access) ซึ่งเป็นสแต็กฮาร์ดแวร์และซอฟต์แวร์ที่แตกต่างกันอย่างสิ้นเชิง ซึ่งช่วยให้เข้าถึงหน่วยความจำโดยตรงระหว่างเครื่องจักรได้โดยไม่ต้องเกี่ยวข้องกับ CPU

AWS มี Elastic Fabric Adapter (EFA) ซึ่งเป็นอินเทอร์เฟซเครือข่ายแบบกำหนดเองที่ใช้โปรโตคอลแบบกำหนดเองของ Amazon ที่เรียกว่า Scalable Reliable Datagram (SRD) ต่างจากเครือข่าย TCP/IP แบบดั้งเดิมที่ข้อมูลต้องถูกคัดลอกหลายครั้งระหว่างยูสเซอร์สเปซ เคอร์เนลสเปซ และบัฟเฟอร์เครือข่าย EFA ที่มี RDMA ช่วยให้สามารถถ่ายโอนข้อมูลโดยตรงระหว่างหน่วยความจำ GPU และการ์ดเครือข่าย โดยข้าม CPU ทั้งหมด

ปรัชญาของการออกแบบเครือข่ายประสิทธิภาพสูง

การสร้างระบบเครือข่ายประสิทธิภาพสูงจำเป็นต้องคิดทบทวนสมมติฐานพื้นฐานหลายประการ:

1. ความเป็นเจ้าของบัฟเฟอร์: ต่างจากซ็อกเก็ตแบบดั้งเดิมที่เคอร์เนลจัดการบัฟเฟอร์เครือข่ายและกำหนดให้ต้องคัดลอกระหว่างยูสเซอร์สเปซและเคอร์เนลสเปซ RDMA กำหนดให้แอปพลิเคชันต้องจัดการบัฟเฟอร์ของตนเอง เมื่อแอปพลิเคชันเริ่มต้นการทำงานของเครือข่าย แอปพลิเคชันจะโอนความเป็นเจ้าของบัฟเฟอร์ไปยังการ์ดเครือข่ายจนกว่าการทำงานจะเสร็จสิ้น ซึ่งช่วยขจัดความจำเป็นในการคัดลอกข้อมูล
2. การลงทะเบียนหน่วยความจำ: แอปพลิเคชันต้องลงทะเบียนภูมิภาคหน่วยความจำกับเคอร์เนลของระบบปฏิบัติการ เคอร์เนลตั้งค่าการแมปเสมือนจริงที่ช่วยให้ CPU, GPU และการ์ดเครือข่ายเข้าใจแอดเดรสเสมือนเดียวกันทั้งหมด การลงทะเบียนนี้เป็นการทำงานเพียงครั้งเดียวที่ช่วยให้สามารถถ่ายโอนข้อมูลแบบ zero-copy ในภายหลังได้
3. ระนาบควบคุมเทียบกับระนาบข้อมูล: เครือข่ายประสิทธิภาพสูงแบ่งการทำงานออกเป็นสองหมวดหมู่:
   1. การทำงานของระนาบควบคุม (เช่น การตั้งค่าการเชื่อมต่อและการลงทะเบียนหน่วยความจำ) จะผ่านเคอร์เนลเพื่อให้แน่ใจว่ามีความปลอดภัย
   2. การทำงานของระนาบข้อมูล (การถ่ายโอนข้อมูลจริง) ข้ามเคอร์เนลเพื่อประสิทธิภาพสูงสุด
4. การรับก่อนการส่ง: ไม่มีบัฟเฟอร์ที่จัดการโดยเคอร์เนล แอปพลิเคชันจะต้องโพสต์การทำงานรับล่วงหน้า โดยระบุว่าควรวางข้อมูลขาเข้าที่ใด นี่คือการเปลี่ยนแปลงพื้นฐานจากโมเดลซ็อกเก็ตที่แอปพลิเคชันสามารถรับข้อมูลได้ตลอดเวลา
5. การเสร็จสิ้นตามการโพล: แทนที่จะรอการแจ้งเตือนของเคอร์เนลผ่านกลไกเช่น epoll แอปพลิเคชันจะโพลคิวความสมบูรณ์ของฮาร์ดแวร์โดยตรง ซึ่งช่วยขจัดค่าใช้จ่ายในการเรียกใช้ระบบและช่วยให้ตอบสนองต่อการทำงานที่เสร็จสมบูรณ์ได้ทันที
6. ความตระหนักรู้ถึงโทโพโลยีฮาร์ดแวร์: การทำความเข้าใจและการปรับ ีประสิทธิภาพสำหรับโทโพโลยีฮาร์ดแวร์เป็นสิ่งสำคัญสำหรับการบรรลุประสิทธิภาพสูงสุด

ทำความเข้าใจโทโพโลยีฮาร์ดแวร์

อินสแตนซ์ AWS p5 มีสถาปัตยกรรมภายในที่ซับซ้อน ตามที่แสดงด้านล่าง แต่ละอินสแตนซ์ประกอบด้วยซ็อกเก็ต CPU สองตัวที่สร้างโหนด NUMA สองโหนด โดยแต่ละโหนด NUMA เชื่อมต่อกับสวิตช์ PCIe สี่ตัว:

ขยาย

ภายใต้สวิตช์ PCIe แต่ละตัว เราจะพบ:

• การ์ดเครือข่าย EFA 100 Gbps สี่ใบ
• NVIDIA H100 GPU หนึ่งตัว
• NVMe SSD หนึ่งตัว

เส้นทางข้อมูลสำหรับการถ่ายโอน TCP/IP และ RDMA แสดงให้เห็นความแตกต่างพื้นฐานในแนวทางของพวกเขา:

ขยาย

ด้วย TCP/IP (ด้านซ้าย) ข้อมูลจะต้องถูกคัดลอกหลายครั้ง:

1. จากหน่วยความจำ GPU ไปยังบัฟเฟอร์แอปพลิเคชันในยูสเซอร์สเปซ
2. จากบัฟเฟอร์ยูสเซอร์สเปซไปยังบัฟเฟอร์เคอร์เนล
3. จากบัฟเฟอร์เคอร์เนลไปยังการ์ดเครือข่าย
4. กระบวนการย้อนกลับเกิดขึ้นที่ฝั่งรับ

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

ในทางตรงกันข้าม RDMA (ด้านขวา) ให้การถ่ายโอนข้อมูลแบบ zero-copy ที่แท้จริง:

1. การ์ดเครือข่ายอ่านโดยตรงจากหน่วยความจำ GPU
2. ข้อมูลเดินทางโดยตรงไปยังการ์ดเครือข่ายระยะไกล
3. การ์ดเครือข่ายระยะไกลเขียนโดยตรงไปยังหน่วยความจำ GPU ปลายทาง
4. แอปพลิเคชันจะต้องตรวจสอบคิวความสมบูรณ์ (CQ) ในยูสเซอร์สเปซเพื่อทราบว่าการถ่ายโอนเสร็จสิ้นเมื่อใด

ด้วย RDMA และการจับคู่ฮาร์ดแวร์ที่เหมาะสม การถ่ายโอนข้อมูลระหว่าง GPU สองตัวจะต้องผ่านสวิตช์ PCIe ท้องถิ่นและเครือข่ายเท่านั้น:

ขยาย

ในทางตรงกันข้าม การถ่ายโอน TCP/IP จะต้องคัดลอกข้อมูลหลายครั้งผ่านหน่วยความจำหลัก ทำให้เกิดความแออัดของบัส PCIe อย่างมาก:

ขยาย

การสร้างด้วย libfabric

เราใช้ libfabric ซึ่งเป็นเฟรมเวิร์กที่มีอินเทอร์เฟซทั่วไปสำหรับบริการแฟบริค การใช้งานของเราใช้การทำงาน RDMA สองประเภท:

• RDMA สองด้าน (SEND/RECV) สำหรับข้อความควบคุมที่มีเมทาดาต้าเกี่ยวกับภูมิภาคหน่วยความจำ (เช่น ออฟเซ็ตและขนาด)
• RDMA WRITE ด้านเดียวสำหรับการถ่ายโอนข้อมูลจริง โดยการดำเนินการ WRITE แต่ละรายการจะจัดการกับส่วนหน่วยความจำที่ต่อเนื่องกันหนึ่งส่วน

เส้นทางสู่ประสิทธิภาพสูงสุด

การพัฒนาของเราก้าวหน้าผ่านหลายขั้นตอน:

1. ใช้งานการถ่ายโอนข้อความทางเดียวพื้นฐานโดยใช้ SEND/RECV
2. ขยายไปสู่การสื่อสารแบบสองทิศทาง
3. เพิ่ม GPUDirect RDMA WRITE สำหรับการถ่ายโอน GPU-GPU โดยตรง
4. ขยายเพื่อจัดการกับการถ่ายโอนพร้อมกันหลายรายการ
5. แนะนำการจัดคิวการทำงานเพื่อความทนทาน
6. บรรลุการใช้งานแบนด์วิดท์ 97.4% บนการ์ดเครือข่ายเดี่ยว

เมื่อปรับขนาดเป็น 32 การ์ดเครือข่าย เราได้นำการปรับ ีประสิทธิภาพที่สำคัญหลายประการมาใช้:

1. การจัดคิวการทำงาน: แทนที่จะส่งการทำงานไปยังการ์ดเครือข่ายโดยตรง เราดูแลรักษาคิวระดับแอปพลิเคชัน ซึ่งช่วยให้มีความทนทานต่อความแออัดของเครือข่ายและทำให้โมเดลการเขียนโปรแกรมง่ายขึ้น
2. Network Warmup: การเตรียมการเชื่อมต่อไว้ล่วงหน้าช่วยปรับปรุงประสิทธิภาพการเริ่มต้นระบบให้ดีขึ้น.
3. Multi-threading: ใช้เธรดเฉพาะสำหรับการทำงานเครือข่ายของ GPU แต่ละตัว.
4. การพินคอร์ CPU: การผูกเธรดกับคอร์ CPU เฉพาะเพื่อหลีกเลี่ยงผลกระทบของ NUMA และการพลาดแคช
5. การชาร์ดสถานะ: ลดการแข่งขันระหว่างเธรด
6. การจัดกลุ่มการทำงาน: การส่งการทำงานหลายรายการร่วมกัน
7. การโพสต์การทำงานแบบเลียรี: การทำงานจะถูกจัดคิวในแอปพลิเคชันก่อน หลังจากโพลคิวความสมบูรณ์แล้ว เราจะพยายามส่งการทำงานที่รอดำเนินการไปยังการ์ดเครือข่าย เพื่อให้มั่นใจว่ามีการใช้ทรัพยากรเครือข่ายอย่างมีประสิทธิภาพ
8. การจัดสรรทรัพยากรที่คำนึงถึง NUMA: จัดสรรทรัพยากร libfabric เช่น คิวความสมบูรณ์บนโหนด NUMA ที่ถูกต้องเพื่อลดความหน่วงในการเข้าถึงหน่วยความจำ

จากการปรับ ีประสิทธิภาพเหล่านี้ เราบรรลุประสิทธิภาพสุดท้ายที่ 3,108 Gbps ซึ่งคิดเป็น 97.1% ของแบนด์วิดท์สูงสุดทางทฤษฎี

วิดีโอด้านล่างแสดงโปรแกรมบรรทัดคำสั่งของเราในการทำงาน โดยจะถ่ายโอนส่วนหน่วยความจำ GPU ที่ไม่ต่อเนื่องไปยังโหนดระยะไกล บรรลุความเร็วในการส่ง 3108.283 Gbps ซึ่งแสดงให้เห็นการใช้แบนด์วิดท์โครงสร้างพื้นฐานเครือข่ายที่เกือบจะเท่ากับทฤษฎี:

บทสรุป

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

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

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

การอ้างอิง

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}