Compiler API

ข้อมูลเบื้องต้น

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

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

วัตถุประสงค์

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

ภาพรวม

• การแยกกลุ่มอินพุต/เอาต์พุต - แนบการแยกกลุ่มกับอินพุตหรือเอาต์พุตของฟังก์ชันหลักเพื่อระบุว่านี่คือวิธีที่ควรแยกกลุ่มเทนเซอร์อินพุต/เอาต์พุตเมื่อส่งไปยัง/แสดงผลจากฟังก์ชัน

• ข้อจำกัดการแยกกลุ่ม - แนบการแยกกลุ่มไปยังเทนเซอร์ระดับกลาง (เช่น ผลลัพธ์ของ matmul) เพื่อระบุว่านี่คือวิธีที่ควรแยกกลุ่มเทนเซอร์นั้นหรือกลุ่มย่อยของการใช้งาน

• กลุ่มการแยกกลุ่ม - จัดกลุ่มเทนเซอร์หลายรายการตามรหัสเพื่อบ่งบอกว่าควรแยกกลุ่มด้วยวิธีเดียวกัน

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

การออกแบบในรายละเอียด

การแยกส่วนอินพุต/เอาต์พุต

อนุญาตให้ผู้ใช้ระบุการแยกข้อมูลสําหรับอินพุตและเอาต์พุตของฟังก์ชันหลัก

ใน MLIR คุณสามารถแนบแอตทริบิวต์กับอาร์กิวเมนต์และผลลัพธ์ของฟังก์ชันได้ ดังนั้นผู้ใช้จึงแนบแอตทริบิวต์การแยกข้อมูลไปยังฟังก์ชันด้วยวิธีนี้ได้

เช่น

@mesh_xy = <["x"=2, "y"=2]>

// The 1st input has a sharding specified, but the 2nd input doesn't.
// The output has a sharding specified.
func @main(%arg0: tensor<8x8xf32>
            {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"}, {}]>},
            %arg1: tensor<8x16xf32>)
    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_xy, [{}, {"y"}]>}) {
  ...
}

ข้อจำกัดการแยกกลุ่ม

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

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

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

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

@mesh_xy = <["x"=2, "y"=2]>

%0 = ... : tensor<8x8xf32>
%1 = sdy.sharding_constraint %0 <@mesh_xy, [{"x"}, {?}]> : tensor<8x8xf32>

กลุ่มการแยกกลุ่ม

ในกรณีที่ไม่มีข้อมูลพึ่งพาหรือข้อมูลพึ่งพาที่แน่นหนาระหว่างเทนเซอร์ 2 รายการขึ้นไป ในขณะที่ผู้ใช้ทราบว่าควรแบ่งเทนเซอร์เหล่านั้นในลักษณะเดียวกันหรือคล้ายกัน Shardy API มีวิธีระบุความสัมพันธ์นี้ ซึ่งช่วยให้ผู้ใช้มีอิสระในการระบุอย่างชัดเจนว่าควรแบ่งกลุ่มเทนเซอร์แยกกัน

เราจึงขอแนะนําแนวคิดกลุ่มกลุ่มย่อย ซึ่งแต่ละกลุ่มจะมีคําสั่งจํานวนเท่าใดก็ได้ที่เชื่อมโยงกับรหัสกลุ่มกลุ่มย่อยเดียวกัน กลุ่มการแยกข้อมูลจะบังคับให้ใช้การแยกข้อมูลภายในกลุ่มเดียวกัน

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

หากเราเรียกใช้โปรแกรมนี้ การนำไปใช้งานการแยกข้อมูลจะไม่สามารถทำนายการแยกข้อมูลของเทมพอร์ล %1 และ %2 และระบบจะทําซ้ำข้อมูลดังกล่าว อย่างไรก็ตาม เราได้แนบแอตทริบิวต์ shard_group ที่ระบุว่าอินพุต %0 และเอาต์พุต %2 อยู่ใน shard_group เดียวกัน ซึ่งจะช่วยให้เราอนุญาตให้มีการนำไปใช้การแยกส่วน @mesh_xy, [{"x"},{"y"}]> จากอินพุต %0 ไปยังเอาต์พุต %2 และไปยังส่วนที่เหลือของกราฟ ซึ่งมีการออกอากาศ %1 คงที่ที่นี่ เรากําหนดค่าให้กับกลุ่มได้ด้วยการดำเนินการ sdy.sharding_group

@mesh_xy = <["x"=2, "y"=2]>

module @"jit_zeros_like" {
  func.func @main(%arg0: tensor<8x2xi64> {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"},{"y"}]>} }) -> (tensor<8x2xi64>) {
    %0 = sdy.sharding_group %arg0, id=0 : tensor<8x2xi64>
    %1 = stablehlo.constant dense<0> : tensor<8x2xi64>
    %2 = sdy.sharding_group %1, id=0 : tensor<8x2xi64>
    return %2 : tensor<8x2xi64>
  }
}

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

การคํานวณด้วยตนเอง

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

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

เช่น

@mesh_name = <["data"=2, "model"=2]>

%0 = ... : tensor<16x32xf32>
%1 = sdy.manual_computation(%0)
    in_shardings=[<@mesh_name, [{"data"}, {"model",?}]>]
    out_shardings=[<@mesh_name, [{"data"}, {?}]>]
    manual_axes={"data"}
    (%arg1: tensor<8x32xf32>) {
  // body
  return %42 : tensor<8x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

อินตัวแปร

1. in_shardings, out_shardings และ manual_axes ทั้งหมดต้องอ้างอิงถึงเมชเดียวกัน manual_axes จัดเรียงตามเมช

2. manual_axes ต้องใช้อย่างชัดเจนในการแยกส่วนข้อมูลเข้า/ออกทั้งหมด กล่าวคือ สำหรับการแยกแต่ละส่วน แกนที่กำหนดเองทั้งหมดต้องแยกมิติข้อมูลหรือทำซ้ำอย่างชัดเจน

3. หากแกนอิสระ (แกนเมชที่ไม่ได้อยู่ใน manual_axes) อยู่ในการจัดกลุ่มอิน/เอาต์รายการใดรายการหนึ่ง แกนนั้นต้องเป็นแกนรองของแกนที่กำหนดเองในการจัดกลุ่มมิติข้อมูลเดียวกัน (ในตัวอย่างนี้ การจัดกลุ่มมิติข้อมูล {"model", "data"} จะใช้งานไม่ได้)

4. ภูมิภาค/เนื้อหาของการคํานวณคือการคํานวณในเครื่อง (เช่น รวมกลุ่มที่ผู้ใช้ระบุ) ต้องเป็นข้อมูลในเครื่องสำหรับการจัดสรรข้อมูลขาเข้า/ขาออก вдоль แกนที่กำหนดเอง (ดูหมายเหตุด้านบน)

การฝังการคํานวณด้วยตนเอง

คุณสามารถฝังการคํานวณด้วยตนเองหลายรายการไว้ด้วยกันได้ ตราบใดที่แต่ละรายการทํางานบนแกนด้วยตนเองชุดที่ไม่ซ้ำกัน