-sdy-close-shardings

ปิดการแบ่งส่วนของ Tensor และนำแกนที่จำลองออก

-sdy-constant-or-scalar-merger

ผสานค่าคงที่และการขยายสเกลาร์ที่เหมือนกันกับการแบ่งส่วนที่ตรงกัน

ทำการเพิ่มประสิทธิภาพการคำนวณย่อย (CSE) แบบเบาในค่าคงที่ที่มีการแบ่งส่วนเหมือนกัน

ไปป์ไลน์การนำเข้าจะแยกและทำซ้ำค่าคงที่และการขยายสเกลาร์เพื่อให้ไม่มีการเผยแพร่การแบ่งส่วนระหว่างการใช้งานที่แตกต่างกันของการคำนวณย่อยแบบค่าคงที่ หากค่าคงที่มีการแบ่งส่วนเหมือนกันหลังจากการเผยแพร่ การส่งผ่านนี้จะผสานค่าคงที่เพื่อประหยัดเวลาในการคอมไพล์ ดูข้อมูลเพิ่มเติมได้ที่ -sdy-constant-or-scalar-splitter

-sdy-convert-global-to-local

แปลงโปรแกรม SDY จากรูปร่างส่วนกลางเป็นรูปร่างในเครื่อง

แปลงโปรแกรม SDY จากรูปร่างส่วนกลางเป็นรูปร่างในเครื่องโดยการแบ่งมิติข้อมูลเชิงตรรกะตามแอตทริบิวต์การแบ่งส่วน

การส่งผ่านนี้ใช้ตัวแปลงประเภทเพื่อแมป RankedTensorType จากรูปร่างเชิงตรรกะส่วนกลางเป็นรูปร่างทางกายภาพในอุปกรณ์

ตัวเลือก

-per-dim-all-gather                     : Keep per-dimension all-gather without combining them into a single all-gather.
-combine-multi-dimension-reduce-scatter : Combine multi-dimension reduce-scatter into a single reduce-scatter.
-enable-rgv3                            : Use StableHLO ReplicaGroupV3 (mesh-axes based) for collectives.
-replica-count                          : Number of replicas (data parallelism).
-partition-count                        : Number of partitions (model parallelism).

-sdy-drop-sharding-and-mesh

นำการดำเนินการ Mesh และสัญกรณ์การแบ่งส่วนออกจากโปรแกรม

-sdy-drop-sharding-rules

นำ OpShardingRuleAttr ออกจากการดำเนินการที่ลงทะเบียนทั้งหมด

-sdy-export-named-computations

สรุปการเรียกจาก NamedComputationOp.

สร้างการส่งผ่านที่แปลง NamedComputationOp เป็น CallOp ด้วยฟังก์ชันส่วนตัวใหม่ที่เรียกว่า name ของ NamedComputationOp FuncOp และ CallOp ใหม่มีการแบ่งส่วนเหมือนกับตัวถูกดำเนินการ/ผลลัพธ์ของ NamedComputationOp เดิม

หากมีฟังก์ชันที่มีชื่อเหมือนกับ NamedComputationOp ในโมดูล ตารางสัญลักษณ์ MLIR จะเปลี่ยนชื่อเป็น {name}_#

-sdy-insert-explicit-reshards

แทรกการแบ่งส่วนใหม่ที่ชัดเจนเพื่อให้การดำเนินการทั้งหมดมีการแบ่งส่วนที่เข้ากันได้

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

หลังจากการเผยแพร่ การดำเนินการบางอย่างอาจยังคงมีการแบ่งส่วนที่ไม่เข้ากัน

โปรดทราบว่าเมื่อใช้แกน (หรือแกนย่อย) เพื่อแบ่งส่วนมิติข้อมูลที่ไม่เกี่ยวข้อง (เช่น มิติข้อมูลที่ไม่หดตัวใน Matmul) ใน Tensor หลายรายการ หรือเมื่อแกนแบ่งส่วนมิติข้อมูลใน Tensor หนึ่งรายการแต่ไม่ได้แบ่งส่วนมิติข้อมูลที่เกี่ยวข้องใน Tensor อื่นๆ เราจะกล่าวว่าการดำเนินการมีการขัดแย้งในการแบ่งส่วน ดังนั้น หลังจากการส่งผ่านนี้ การดำเนินการจะไม่มีการขัดแย้ง

การส่งผ่านนี้จะแทรกการดำเนินการแบ่งส่วนใหม่อย่างชัดเจนเพื่อให้มิติข้อมูลที่เกี่ยวข้องมีการแบ่งส่วนในลักษณะเดียวกันในตัวถูกดำเนินการและผลลัพธ์ทั้งหมด และแกน (หรือแกนย่อย) แต่ละแกนจะใช้เพื่อแบ่งส่วนมิติข้อมูลประเภทเดียวเท่านั้น

ตัวอย่าง

อินพุต:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

เอาต์พุต:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

ในตัวอย่างข้างต้น lhs และ rhs มีการแบ่งส่วนในแกน "x" ในมิติข้อมูลที่ไม่หดตัว ซึ่งไม่เข้ากัน การส่งผ่านจะแทรกการแบ่งส่วนใหม่ที่ชัดเจนใน rhs ก่อนการดำเนินการ Dot เพื่อให้การดำเนินการ Dot มีการแบ่งส่วนที่เข้ากันได้

ตัวเลือก

-enable-full-version : Enable full version.

-sdy-insert-func-call-reshards

แทรกการแบ่งส่วนใหม่สำหรับการขัดแย้งในการแบ่งส่วน Func และ Call

แทรกการแบ่งส่วนใหม่สำหรับการขัดแย้งในการแบ่งส่วน Func และ Call ในผลลัพธ์

-sdy-optimize-collectives

เพิ่มประสิทธิภาพการดำเนินการสื่อสารแบบกลุ่ม

ขจัดการดำเนินการ Collective-Permute ที่ซ้ำซ้อนซึ่งนำหน้าเชน All-to-All โดยการแยกมิติข้อมูลที่แยกออกและผสานรวมเป็น All-to-All แบบรวม

-sdy-pad-for-divisibility

เพิ่ม Tensor ที่มีการแบ่งส่วนที่ไม่สามารถหารได้ลงในรูปร่างที่สามารถหารได้

-sdy-per-instruction-partitioning

แยกคำสั่งเป้าหมายออกอย่างเลือกสรร เรียกใช้ไปป์ไลน์ตัวแบ่งส่วนในคำสั่งเหล่านั้นเพื่อสร้างโค้ดในอุปกรณ์ และห่อโค้ดเหล่านั้นใน sdy.manualcomputation

สำหรับการแก้ไขข้อบกพร่องและการแบ่งส่วนแบบครึ่งต่อครึ่ง ให้แบ่งส่วนคำสั่งแต่ละรายการอย่างเลือกสรรโดยใช้ไปป์ไลน์ตัวแบ่งส่วน Shardy แบบสแตนด์อโลน (resolve-permutation-factors, reshard-to-collectives, pad-for-divisibility, convert-global-to-local) และห่อโค้ดในอุปกรณ์ที่ได้ไว้ใน sdy.manual_computation ซึ่งจะทำให้โปรแกรมส่วนกลางที่เหลืออยู่เหมือนเดิม เพื่อให้การแบ่งส่วน XLA SPMD แบ่งส่วนเฉพาะคำสั่งส่วนกลางที่ไม่ได้ห่อ

ตัวเลือก

-filter               : Filter string for selective partitioning. Can be empty (all sharded ops), comma-separated op name substrings (e.g. 'dot, pad'), or key-values like 'selectLow=0, selectHigh=10'.
-enable-halo-exchange : Implement halo exchange logic for windowed operations inside the sub-pipeline.
-replica-count        : Number of replicas (data parallelism).
-partition-count      : Number of partitions (model parallelism).

-sdy-propagate-to-func-results

เผยแพร่การแบ่งส่วนจากตัวสิ้นสุด Func ไปยังผลลัพธ์ Func

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

-sdy-remove-all-gather-reduce-scatter-for-cmv1

นำ sdy.all_gather และ sdy.reducescatter ออกสำหรับ CMV1

นำ All-Gather ออกในรูปแบบ All-Gather + Dot นำ Reduce-Scatter ออกในรูปแบบ Dot + Reduce-Scatter การส่งผ่านนี้มีไว้เพื่อความเข้ากันได้กับ Collective Matmul V1 (CMV1) ซึ่งเป็นโซลูชันชั่วคราวสำหรับ b/432019089

-sdy-remove-propagation-debug-info

นำข้อมูลการแก้ไขข้อบกพร่องการเผยแพร่ (ขอบการเผยแพร่และการแบ่งส่วนต้นทาง) ออกระหว่างการส่งออก

-sdy-remove-sharding-groups

นำ ShardingGroupOps ออกหลังจากการเผยแพร่

-sdy-remove-sub-axes-in-input-output-shardings

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

ผู้ใช้ Shardy บางรายคาดหวังว่าอินพุต/เอาต์พุตของฟังก์ชันจะมีการแบ่งส่วนโดยไม่มีแกนย่อย การส่งผ่านนี้จะนำแกนย่อยและแกนต่อท้ายออกจากมิติข้อมูลการแบ่งส่วนอินพุต/เอาต์พุตแบบเปิด โดยปกติการส่งผ่านนี้จะเกิดขึ้นหลังจากการส่งผ่าน sdy-update-non-divisible-input-output-shardings เพื่อให้แน่ใจว่าการนำแกนย่อยออกจะไม่ทำให้เกิดการแบ่งส่วนที่ไม่สามารถหารได้

-sdy-reshard-to-collectives

แปลง ReshardOp เป็นการดำเนินการแบบกลุ่ม Shardy ต่างๆ

จับคู่การดำเนินการ Reshard และเขียนใหม่เป็นการดำเนินการแบบกลุ่ม Shardy ต่างๆ หลังจากการส่งผ่านนี้ จะไม่มีการดำเนินการ Reshard เหลืออยู่ในโมดูล

หาก keepRedundantReshards เป็นจริง การดำเนินการ Reshard ที่เหลืออยู่จะเป็นการดำเนินการที่ซ้ำซ้อนเท่านั้น โดยค่าเริ่มต้น ระบบจะถือว่ามีการแทรกการแบ่งส่วนใหม่ที่ชัดเจนแล้ว (sdy-insert-explicit-reshards) และจะไม่เก็บการแบ่งส่วนใหม่ที่ซ้ำซ้อน ระบบควรเก็บการแบ่งส่วนใหม่ที่ซ้ำซ้อนไว้หากยังไม่ได้แทรกการแบ่งส่วนใหม่ที่ชัดเจน

ตัวอย่าง

อินพุต:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

เอาต์พุต:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.all_gather \[{"y", "z"}, {}\] %arg0 out_sharding=<@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

ในตัวอย่างข้างต้น Tensor %0 : tensor<16x2xf32> มีการแบ่งส่วนเป็น \[{"x", "y", "z"}, {}\] จากนั้นมีการดำเนินการ reshard ที่แบ่งส่วนใหม่เป็น \[{"x"}, {}\] ในแกนแรก เนื่องจากมีการนำคำต่อท้าย {"y", "z"} ออก หลังจากการแบ่งส่วนใหม่ เราจึงอนุมานได้ว่าเราได้รวบรวม {"y", "z"} ทั้งหมด มิติข้อมูลที่ 2 ไม่มีการเปลี่ยนแปลง

ตัวเลือก

-keep-redundant-reshards : Whether it keeps redundant reshards or removes.

-sdy-resolve-permutation-factors

แก้ปัญหาการแบ่งส่วนในมิติข้อมูลที่แมปกับปัจจัย kPermutation

มิติข้อมูลการแบ่งส่วนที่มีปัจจัย kPermutation อาจต้องมีการสื่อสารข้ามอุปกรณ์ (เช่น การแลกเปลี่ยนรัศมีสำหรับหน้าต่างหรือการสลับแบบกลุ่มสำหรับการย้อนกลับ)

หาก enableHaloExchange เป็นจริง การส่งผ่านจะใช้ตรรกะการสื่อสารที่เพิ่มประสิทธิภาพแล้วที่มีอยู่เพื่อแก้ปัญหาปัจจัยการสลับ มิเช่นนั้น การส่งผ่านจะแทรกการดำเนินการ sdy.reshard เพื่อจำลองมิติข้อมูลเหล่านั้น ค่าเริ่มต้นของ enableHaloExchange คือจริง

ตัวเลือก

-enable-halo-exchange : Implement halo exchange logic for windowed operations.
-replica-count        : Number of replicas (data parallelism).
-partition-count      : Number of partitions (model parallelism).

-sdy-resolve-single-device-sharding

ลดการดำเนินการในอุปกรณ์เดียวลงใน Stablehlo.if ที่มีการป้องกันด้วยรหัสอุปกรณ์เป้าหมาย

ตัวเลือก

-replica-count   : Number of replicas per partition.
-partition-count : Number of partitions per replica.

-sdy-sharding-constraint-to-reshard

แปลง ShardingConstraintOp เป็น ReshardOp

-sdy-sink-data-flow-edges

รวม DataFlowEdgeOp ทั้งหมดไว้ในอินพุต

ย้ายการแบ่งส่วนของ DataFlowEdgeOp แต่ละรายการไปยังอินพุต (เป้าหมายรากของขอบ) และแทนที่การดำเนินการด้วยอินพุต

ตัวเลือก

-sink-debug-sharding-origins          : Whether to sink the debug sharding origins info. See `debug-sharding-origins` option in propagation for more info.
-sink-debug-propagation-edge-sharding : Whether to sink the debug propagation edge sharding info. See `debug-propagation-edge-sharding` option in propagation for more info.

-sdy-sink-func-data-flow-edges

รวม FuncDataFlowEdgeOp ทั้งหมดไว้ในอินพุต

ย้ายการแบ่งส่วนของ FuncDataFlowEdgeOp แต่ละรายการไปยังอินพุตและแทนที่การดำเนินการด้วยอินพุต

-sdy-unflatten-call-graph

ยกเลิกการทำให้กราฟการเรียกแบน

ยกเลิกการทำให้กราฟแบน โดยจะขจัดฟังก์ชันที่ซ้ำกันซึ่งมีการแบ่งส่วนอินพุต/เอาต์พุตเหมือนกัน และมีต้นทางเดียวกันตามที่อธิบายไว้ในแอตทริบิวต์ 'original_func_name' ที่แนบมากับฟังก์ชัน

ตัวเลือก

-dedup-functions-fully : If true, regardless of the input and output shardings of functions, it keeps one callee function for each caller function. The default is false, meaning it will deduplicate only if the input and output shardings are the same.

-sdy-update-non-divisible-input-output-shardings

ทำให้อินพุต/เอาต์พุต FuncOp มีการแบ่งส่วนอย่างเท่าเทียมกัน โดยไม่จำเป็นต้องเพิ่มเนื่องจากการแบ่งส่วนที่ไม่สามารถหารได้

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

-sdy-verify-unreduced-axes

ตรวจสอบความสอดคล้องของการใช้แกนที่ไม่ได้ลด

ตรวจสอบว่าสำหรับการดำเนินการทุกรายการ หากตัวถูกดำเนินการมีแกนที่ไม่ได้ลด การดำเนินการจะลดแกนเหล่านั้นอย่างชัดเจน (เช่น ผ่าน sdy.reshard) หรือส่งแกนเหล่านั้นไปยังผลลัพธ์ (หรือเป็นการดำเนินการแบบขอบเขต เช่น func.call)