การจัดสรรที่ไม่คำนึงถึงรูปแบบ

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

กฎการแยกข้อมูล

กฎการแยกข้อมูลจะเข้ารหัสวิธีที่เราเผยแพร่ผ่านการดำเนินการ เนื่องจากตอนนี้ Shardy ขึ้นอยู่กับ StableHLO จึงกำหนดกฎการแยกส่วนสำหรับการดำเนินการ StableHLO แต่ละรายการ นอกจากนี้ Shardy ยังมี ShardingRuleOpInterface ซึ่งเจ้าของภาษาถิ่นสามารถใช้ในการดำเนินการเพื่อกำหนดกฎการแยกส่วนสำหรับการดำเนินการของตนเอง ตราบใดที่การดำเนินการใช้อินเทอร์เฟซนี้ Shardy จะนำไปใช้งานได้

def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
  let methods = [
    InterfaceMethod<
      /*desc=*/[{
        Returns the sharding rule of the op.
      }],
      /*retType=*/"mlir::sdy::OpShardingRuleAttr",
      /*methodName=*/"getShardingRule"
    >,
  ];
}

การดำเนินการกับโฟลว์ข้อมูล

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

def ShardableDataFlowOpInterface :
    OpInterface<"ShardableDataFlowOpInterface"> {
  (get|set)BlockArgumentEdgeOwnerShardings;
  (get|set)OpResultEdgeOwnerShardings;
  getBlockArgumentEdgeOwners;
  getOpResultEdgeOwners;
  getEdgeSources;
  // ...
}

ดูภาพรวมระดับสูงเกี่ยวกับวิธีที่เราจัดการการดำเนินการกับข้อมูลได้ที่การดำเนินการกับข้อมูล

อินเทอร์เฟซที่ยังไม่ได้ใช้งาน

ในอนาคต เราจะเพิ่มอินเทอร์เฟซและลักษณะอื่นๆ เพื่อให้ Shardy มีความยืดหยุ่นมากขึ้นและไม่ขึ้นอยู่กับภาษา โปรดดูรายการด้านล่าง

การแยกแบบคงที่

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

ตัวอย่างเช่น ในรูปภาพด้านล่างนี้ หากมีการแบ่ง add ออกเป็นหลายกลุ่ม การดำเนินการนี้จะไม่ส่งผลต่อวิธีแบ่ง divide และ subtract (ในส่วนต่างๆ ของการประมวลผล) ออกเป็นหลายกลุ่ม

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

ผู้ใช้ Shardy จะต้องกำหนดสิ่งต่อไปนี้เพื่อให้บรรลุเป้าหมายนี้ - พาส your_dialect.constant -> sdy.constant - แอตทริบิวต์ sdy::ConstantLike เช่น iota - แอตทริบิวต์ mlir::Elementwise สำหรับการดำเนินการแบบองค์ประกอบ เช่น add และ multiply - sdy::ConstantFoldable สำหรับการดำเนินการ เช่น slice/broadcast ในทางเทคนิคแล้ว ระบบจะคํานวณการดำเนินการเหล่านี้ได้เมื่อคอมไพล์ หากโอเปอเรนด์/ผลลัพธ์ทั้งหมดเป็นค่าคงที่

ลำดับความสำคัญของการดำเนินการ

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

รูปภาพด้านล่างแสดงวิธีใช้ลําดับความสําคัญใน GSPMD เพื่อเผยแพร่การดำเนินการตามลําดับที่ถูกต้อง

ดูการอภิปรายเกี่ยวกับเหตุผลที่ความสำคัญสูงสุดมีความสําคัญได้จากเอกสาร GSMPD

ทำงานได้กับทุกภาษา

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