Shardy เป็นระบบการแบ่งกลุ่มเทนเซอร์ที่ใช้ MLIR สำหรับทุกภาษา เครื่องมือนี้สร้างขึ้นจากความร่วมมือของทีม GSPMD และ PartIR โดยรวมเอาข้อดีของทั้ง 2 ระบบและประสบการณ์ที่ทีมและผู้ใช้ได้แชร์ไว้
ประโยชน์
- ผู้ใช้สามารถควบคุมและคาดการณ์ได้มากขึ้นด้วยการรวมการนำไปใช้งานของ GSPMD กับการแบ่งพาร์ติชันแบบเพิ่มของ PartIR
- ฟีเจอร์ใหม่ที่ขับเคลื่อนโดยประสบการณ์การใช้งาน เช่น การรองรับการรีเชปรูปแบบใหม่ซึ่งมักทำให้เกิดการสื่อสารเพิ่มเติม เว้นแต่ผู้ใช้จะทราบวิธีดำเนินการ
- ความสามารถในการใช้งานและการแก้ไขข้อบกพร่องที่ดีขึ้นเพื่อเพิ่มความเร็วของผู้ใช้ปลายทาง เช่น การใช้การแสดงภาพการแยกข้อมูลตามแกน
- โค้ดเบสแบบโอเพนซอร์สที่เรียบง่ายซึ่งใช้ MLIR โดยมีผู้มีส่วนร่วมที่ใช้งานอยู่จำนวนมากขึ้น (ภายใน ภายนอก และจากเขตเวลาต่างๆ) เพื่อสนับสนุนผู้ใช้
คอมโพเนนต์
- การนำเสนอการแยกส่วน: การนำเสนอการแยกส่วนที่อิงตามแกนซึ่งเชื่อมโยงกับเมชเชิงตรรกะหนึ่งๆ (จากเมชหลายรายการที่เป็นไปได้) และรองรับการจำกัดการแยกส่วนและแกนของมิติข้อมูล การแยกแกนสำหรับการดำเนินการต่างๆ เช่น การปรับรูปร่างใหม่ ลำดับความสำคัญสำหรับการแบ่งพาร์ติชันแบบเพิ่ม และอื่นๆ
- Compiler API: ชุดคอมโพเนนต์คอมไพเลอร์ที่ใช้ร่วมกับการนำเสนอการแยกส่วนเพื่อส่งผลต่อการนำไปใช้การแยกส่วน
- การแยกกลุ่มอินพุต/เอาต์พุต - แนบการแยกกลุ่มกับอินพุตหรือเอาต์พุตของฟังก์ชันหลักเพื่อระบุว่านี่คือวิธีที่ควรแยกกลุ่มเทนเซอร์อินพุต/เอาต์พุตเมื่อส่งไปยัง/แสดงผลจากฟังก์ชัน
- ข้อจำกัดการแยกกลุ่ม - แนบการแยกกลุ่มไปยังเทนเซอร์ระดับกลาง (เช่น ผลลัพธ์ของ matmul) เพื่อระบุว่าควรแยกกลุ่มเทนเซอร์นั้นหรือกลุ่มย่อยของการใช้งาน
- แยกเป็น/เหมือน - จัดกลุ่มเทนเซอร์หลายรายการตามรหัสเพื่อระบุว่าควรแยกเป็นกลุ่มด้วยวิธีเดียวกัน
- การคํานวณด้วยตนเอง - ล้อมรอบการคํานวณย่อยที่มีการแบ่งพาร์ติชันด้วยตนเองโดยใช้ชุดย่อยของแกนเมช โดยมีการระบุการแยกกลุ่มตามแกนด้วยตนเองเหล่านั้นสําหรับอินพุตและเอาต์พุตทั้งหมด และภายในการคํานวณย่อย ประเภทเทนเซอร์จะเป็นแบบภายในเมื่อเทียบกับการแยกกลุ่มเหล่านั้น
- การนำไปใช้งานการแยกกลุ่ม: อัลกอริทึมการนำไปใช้งานซึ่งรวมลำดับความสำคัญของผู้ใช้และข้อจำกัดการแยกกลุ่มเข้ากับรูปแบบต้นทุนและวิธีการแก้ปัญหาแบบเฮuristic ของคอมไพเลอร์ ดังนี้
- ลำดับความสำคัญที่ผู้ใช้กำหนด เช่น ทำแบบขนานของกลุ่มแล้วทำ ZeRO
- ลำดับความสำคัญตามการดำเนินการ เช่น การดำเนินการแบบองค์ประกอบก่อนแล้วตามด้วย matmul เป็นต้น
- วิธีการแก้ปัญหาแบบละเอียดยิ่งขึ้น เช่น ชอบมิติข้อมูลการประมวลผลเป็นกลุ่ม
- ตัวแบ่งพาร์ติชัน SPMD: คอมโพเนนต์ที่ลดการตัดสินใจการนำไปใช้งานการแยกส่วนโดยการแบ่งพาร์ติชันโปรแกรมออกเป็นโปรแกรม SPMD เพิ่มการเคลื่อนไหว/การจัดรูปแบบข้อมูลที่จำเป็น และการดำเนินการแบบรวมในกระบวนการ
- ในระยะสั้น การติดตั้งใช้งานครั้งแรกจะใช้โปรแกรมแบ่งพาร์ติชัน SPMD ของ GSPMD ปัจจุบัน
- ในระยะยาว เราวางแผนที่จะสร้างโปรแกรมแบ่งกลุ่ม SPMD ที่ใช้ MLIR ใหม่
ที่เก็บโค้ด
โปรเจ็กต์ Shardy กำลังอยู่ในระหว่างการพัฒนา และเราต้องการความคิดเห็นจากชุมชนโอเพนซอร์ส โค้ด Shardy มีอยู่ที่ https://github.com/openxla/shardy