คู่มือปฏิบัติ: การสื่อสารที่เริ่มต้นโดย GPU เพื่อการจำลองพลวัตโมเลกุลในระดับสเกล
การจำลองพลวัตโมเลกุล (Molecular Dynamics - MD) ถือเป็นหนึ่งในภาระงานที่ต้องการทรัพยากรในการประมวลผลสูงที่สุดในวิทยาศาสตร์เชิงคำนวณ ด้วยการจำลองเหล่านี้ นักวิจัยสามารถสังเกตพฤติกรรมของอะตอมได้อย่างละเอียด ตั้งแต่การพับตัวของโปรตีนไปจนถึงการค้นคว้าหายาและวัสดุใหม่ๆ
อย่างไรก็ตาม การจำลองเหล่านี้ต้องแลกมาด้วยต้นทุนที่สูง โดยทั่วไปแล้ว การจำลองจะครอบคลุมอะตอมหลายแสนถึงหลายสิบล้านอะตอม และต้องก้าวไปข้างหน้าทีละเฟมโตวินาที ผ่านการคำนวณหลายพันล้านขั้น เนื่องจากขนาดของปัญหาที่มักจะคงที่ การขนาน (parallelization) จึงต้องกระจายงานไปยังทรัพยากรที่มีอยู่ทั้งหมดพร้อมกัน ซึ่งเป็นสภาวะที่เรียกว่า "การขยายสเกลแบบแข็ง" (strong scaling)
GROMACS เป็นหนึ่งในแพ็กเกจ MD ที่ใช้กันอย่างแพร่หลายที่สุดในโลก ฮาร์ดแวร์สมัยใหม่และการประมวลผลแบบขนานบน CPU-GPU ได้ผลักดันประสิทธิภาพไปสู่ระดับต่ำกว่ามิลลิวินาที โดยสามารถประมวลผลได้ 100–200 ไมโครวินาทีต่อขั้นเวลาในหลายๆ GPU การบรรลุ strong scaling ในระดับนี้ต้องการความหน่วงแฝง (latency) ของเคอร์เนลที่ต่ำมาก แต่การสื่อสารจาก GPU ไปยัง GPU ยังคงเป็นคอขวดพื้นฐานเมื่อการจำลองขยายสเกลไปยัง GPU จำนวนมากขึ้น
ข้อจำกัดของการสื่อสารแบบดั้งเดิมใน GROMACS
แอปพลิเคชัน HPC ขนาดใหญ่ส่วนใหญ่ รวมถึง GROMACS ใช้ Message Passing Interface (MPI) สำหรับการสื่อสารระหว่างกระบวนการ (inter-process communication) อย่างไรก็ตาม MPI ถูกออกแบบมาสำหรับการประมวลผลที่เน้น CPU เป็นหลัก
เมื่อ GROMACS ทำงานบน GPU กระบวนการสื่อสารจะบังคับให้ GPU ต้องหยุดทำงานในขณะที่ CPU จัดการการถ่ายโอนข้อมูลก่อนที่จะส่งสัญญาณให้ GPU ทำงานต่อ ในส่วนของ GROMACS ที่เรียกว่า "halo exchange" ซึ่งเป็นการแบ่งปันข้อมูลอะตอมบริเวณขอบระหว่างโดเมน GPU ที่อยู่ติดกัน กระบวนการนี้จะเกิดขึ้นซ้ำในทั้งสามมิติเชิงพื้นที่ (3D decompositions) ซึ่งใช้เวลา CPU มากกว่า 50% ของเวลาทั้งหมดในการทำงานในอัตราสูงสุด และจำกัดความสามารถในการขยายสเกล
การเร่งความเร็ว GROMACS ด้วยการเข้าถึงหน่วยความจำระยะไกลที่เริ่มต้นโดยอุปกรณ์ (Device-initiated Remote Memory Access)
กุญแจสำคัญในการเร่งความเร็ว GROMACS คือการขจัดขั้นตอนการส่งมอบงานระหว่าง CPU และ GPU ด้วยการสื่อสารแบบ Native ของ GPU NVIDIA NVSHMEM เป็นไลบรารีสำหรับการใช้งาน remote memory access โดยอิงตามโมเดล partitioned global address space ของ OpenSHMEM NVSHMEM ช่วยให้เคอร์เนล GPU สามารถเริ่มต้นการถ่ายโอนข้อมูลได้โดยตรง ลดการพึ่งพา CPU ในเส้นทางวิกฤต และช่วยให้การทับซ้อนของการสื่อสารและการคำนวณมีประสิทธิภาพมากขึ้น
GROMACS halo exchange ใช้กลไกการส่งต่อข้อมูลเป็นระยะ (staged forwarding mechanism) โดยข้อมูลจะถูกส่งผ่านอันดับ (ranks) ที่อยู่ระหว่างกลางในแต่ละมิติเชิงพื้นที่ ผ่านขั้นตอนการสื่อสารหนึ่งขั้นหรือมากกว่า หรือที่เรียกว่า "pulses" แต่ละ pulse จะเริ่มต้นด้วยเคอร์เนล "pack" ที่รวบรวมอะตอมขอบไปยังบัฟเฟอร์ส่งข้อมูลที่ต่อเนื่องกัน สิ่งนี้สร้างการพึ่งพา (dependency) ระหว่างขั้นตอนการสื่อสาร (Z → Y → X) ซึ่งการใช้งานแบบดั้งเดิมจะจัดการผ่านการทำให้ขั้นตอนต่างๆ เป็นลำดับ (coarse phase-level serialization) การปรับปรุงใหม่นี้จะแทนที่การกั้นแบบหยาบ (coarse barriers) ด้วยสัญญาณต่อ pulse แบบละเอียด (fine-grained per-pulse signals) ที่แสดงผลโดยตรงในเคอร์เนลที่หลอมรวม (fused kernel)
การปรับปรุงประสิทธิภาพการถ่ายโอนข้อมูลระหว่าง GPU ที่เชื่อมต่อด้วย NVLink
ในการออกแบบที่กล่าวมา ข้อมูลจะถูกส่งผ่าน NVSHMEM transport เสมอ (nvshmemxfloatputsignalnbi_block) ซึ่งใช้ได้กับทุกการเชื่อมต่อ แต่เมื่อ GPU ปลายทางเชื่อมต่อด้วย NVLink GPU สามารถเข้าถึงหน่วยความจำของกันและกันได้โดยตรง แทนที่จะแพ็คข้อมูลลงในบัฟเฟอร์ภายในเครื่องและเรียกใช้คำสั่ง put ให้แพ็คข้อมูลลงในอาร์เรย์พิกัดของ GPU ปลายทางโดยตรง และขจัดรอบการเดินทางของหน่วยความจำส่วนกลางที่แยกออกไป
nvshmem_ptr(remotePtr, peerRank) จะคืนค่าพอยน์เตอร์อุปกรณ์ที่ไม่เป็น null เมื่อ GPU ปลายทางสามารถเข้าถึงได้ผ่าน NVLink และคืนค่า null หากไม่สามารถเข้าถึงได้ การตรวจสอบนี้จะทำเพียงครั้งเดียวในเคอร์เนลและแยกการทำงานตามเงื่อนไข
ในเส้นทาง NVLink การวนลูปการแพ็คจะเขียนข้อมูลลงในตำแหน่งที่แน่นอนในอาร์เรย์พิกัดของ GPU ปลายทาง เมื่อบล็อก 0 ยืนยันว่า CTA ทั้งหมดแพ็คข้อมูลเสร็จสิ้นแล้วโดยใช้ device-scoped barrier จะส่งสัญญาณไปยัง GPU ปลายทางด้วย system-scope release store (st.release.sys) ซึ่งรับประกันว่าการเขียนโดยตรงก่อนหน้านี้ทั้งหมดจะมองเห็นได้โดย GPU ปลายทางก่อนที่จะอ่านสัญญาณ ไม่จำเป็นต้องมีคำสั่ง put ข้อมูลเพิ่มเติม เนื่องจากข้อมูลอยู่ในตำแหน่งที่ถูกต้องแล้ว
เมื่อ nvshmemptr คืนค่า null พฤติกรรมจะเหมือนกับการออกแบบแรก ข้อมูลจะถูกแพ็คลงในบัฟเฟอร์พักข้อมูลภายในเครื่อง (local staging buffer) และส่งต่อไปยัง nvshmemxfloatputsignalnbiblock ซึ่ง NVSHMEM transport จะส่งผ่านเครือข่าย RDMA เช่น InfiniBand, Slingshot หรืออื่นๆ
การปรับปรุงนี้เน้นให้เห็นถึงความไม่สมมาตรที่สำคัญระหว่างสองเส้นทาง เส้นทาง NVLink ต้องการ CUDA core จำนวนมาก (SMs) ที่ออกคำสั่ง store พร้อมกันเพื่อใช้ประโยชน์จากแบนด์วิดท์ NVLink ให้เต็มที่ ในขณะที่เส้นทาง InfiniBand ต้องการเพียงเธรดเดียว (CTA) เพื่อเริ่มต้นการถ่ายโอนข้อมูลจำนวนมากโดยใช้ NVSHMEM transport คอขวดที่เหลือคือค่าใช้จ่ายในการเปิดใช้งานเคอร์เนลต่อ pulse และลำดับการทำงานที่ตายตัวระหว่าง pulses แม้ว่าข้อมูลบางส่วนจะไม่มีการพึ่งพาระหว่าง pulses ก็ตาม
การเปิดใช้งานการประมวลผลแบบขนานระหว่าง Pulses ผ่าน Kernel Fusion
แนวคิดหลักคือข้อมูลทั้งหมดใน pulse ไม่ได้ขึ้นอยู่กับ pulse ก่อนหน้าเสมอไป อะตอมส่วนใหญ่ใน pulse สามารถถูกแพ็คและส่งออกได้ทันที มีเพียงอะตอมที่ "ส่งต่อ" (forwarded atoms) ที่ได้รับข้อมูลจาก pulse ก่อนหน้าเท่านั้นที่ต้องรอ สิ่งนี้ทำได้โดยการแบ่งแผนที่ดัชนี (index map) ของแต่ละ pulse ที่ dependencyAtomOffset:
- อะตอมอิสระ (Independent atoms) (atomIndex < dependencyAtomOffset): ปลอดภัยที่จะแพ็คและส่งทันทีที่ pulse เริ่มทำงาน
- อะตอมที่ต้องพึ่งพา (Dependent atoms) (atomIndex ≥ dependencyAtomOffset): ต้องรอสัญญาณ pulse ก่อนหน้าที่มีความเกี่ยวข้อง
ด้วยการแบ่งส่วนนี้ เคอร์เนลเดียวสามารถประมวลผลทุก pulse พร้อมกันได้ งานที่ต้องพึ่งพาจะรอเพียงสัญญาณต่อ pulse ที่ต้องการเท่านั้น ไม่ใช่รอทั้งเฟส
การทดสอบประสิทธิภาพ
การทดสอบบนซูเปอร์คอมพิวเตอร์ NVIDIA Eos และคลัสเตอร์ NVIDIA GB200 NVL72 แสดงให้เห็นถึงประสิทธิภาพการขยายสเกลแบบแข็ง (strong scaling) ที่ดีขึ้นถึง 2 เท่าเมื่อเทียบกับ GPU-aware MPI โดยเฉพาะอย่างยิ่งสำหรับระบบที่จำกัดด้วยความหน่วงแฝง (latency-bound systems) แนวทางนี้สามารถนำไปปรับใช้กับแอปพลิเคชัน HPC ใดๆ ที่มีรูปแบบ halo exchange ได้ แม้ว่าการกำหนดมาตรฐานของ primitive การสื่อสารแบบ Native ของ GPU ยังคงเป็นความท้าทายที่เปิดกว้าง
สรุป
การนำ GPU-initiated communication และ NVSHMEM มาใช้สามารถขจัดคอขวดที่เกิดจากการสื่อสารที่ควบคุมโดย CPU ใน GROMACS ได้อย่างมีประสิทธิภาพ ช่วยเพิ่มประสิทธิภาพการจำลองพลวัตโมเลกุลในระดับสเกล และเปิดโอกาสให้นักวิจัยสามารถสำรวจพฤติกรรมของอะตอมได้อย่างลึกซึ้งยิ่งขึ้น
#GPU #NVSHMEM #MolecularDynamics #HPC #GROMACS
ขอบคุณ แหล่งข้อมูล
https://developer.nvidia.com/blog/a-practical-guide-to-gpu-initiated-communication-for-molecular-dynamics-at-scale/