-sdy-close-shardings
Đóng các phân đoạn tensor và thả các trục được sao chép.
-sdy-drop-sharding-rules
Xoá OpShardingRuleAttr
khỏi tất cả các thao tác đã đăng ký.
-sdy-insert-explicit-reshards
Chèn các phân đoạn lại rõ ràng để tất cả các thao tác đều có các phân đoạn tương thích.
Về cơ bản, việc phân đoạn tương thích có nghĩa là toán tử có thể chấp nhận các toán hạng được phân đoạn và tạo ra kết quả được phân đoạn mà không cần bất kỳ thông tin liên lạc phân đoạn lại nào (lưu ý rằng toán tử vẫn có thể yêu cầu thông tin liên lạc như tất cả các phép giảm hoặc hoán đổi hào quang).
Sau khi truyền tải, một số thao tác vẫn có thể có các phân đoạn không tương thích.
Xin lưu ý rằng khi một trục (hoặc trục phụ) được dùng để phân đoạn các phương diện không tương ứng (ví dụ: các phương diện không co rút trong matmul) trên nhiều tensor, hoặc khi một trục phân đoạn một phương diện trong một tensor nhưng không phải là phương diện tương ứng trong tensor khác, thì thao tác đó được cho là có xung đột phân đoạn. Do đó, sau khi chuyển đổi này, các thao tác sẽ không còn xung đột.
Lệnh truyền này chèn các thao tác phân đoạn lại một cách rõ ràng để đối với mỗi thao tác, các phương diện tương ứng được phân đoạn theo cùng một cách trên tất cả các toán hạng và kết quả, đồng thời mỗi trục (hoặc trục phụ) chỉ có thể được dùng để phân đoạn một loại phương diện.
Ví dụ minh hoạ:
Đầu vào:
mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"y"},{"x"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
: (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
Kết quả:
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>
Trong ví dụ trên, có xung đột vì tensor lhs
và rhs
đều được phân đoạn theo trục "x" trên các phương diện không co giãn. Ở đây, tensor rhs
được phân đoạn lại, trước phép toán dấu chấm, chỉ được phân đoạn rõ ràng trên phương diện đầu tiên và trên trục "x". Bằng cách này, toán tử dấu chấm sẽ tương thích.
-sdy-remove-sharding-groups
Xoá ShardingGroupOps sau khi truyền.
-sdy-sharding-constraint-to-reshard
Chuyển đổi ShardingConstraintOp thành ReshardOp.
-sdy-sink-data-flow-edges
Chuyển tất cả DataFlowEdgeOp
vào đầu vào của chúng.
Di chuyển hoạt động phân đoạn của mỗi DataFlowEdgeOp
đến đầu vào (mục tiêu gốc của cạnh) và thay thế toán tử bằng đầu vào.
TODO(tomnatan): cân nhắc việc di chuyển tính năng phân đoạn sang tất cả các mục tiêu có thể có tính năng phân đoạn đính kèm.
-sdy-update-non-divisible-input-output-shardings
Phân đoạn đầu vào/đầu ra FuncOp một cách đồng đều, không cần phải thêm khoảng đệm do các phân đoạn không thể chia.
Người dùng Shardy mong muốn đầu vào/đầu ra của hàm có thể chia đều/phân đoạn để tránh yêu cầu thêm vào các tensor. Quá trình truyền tải có thể khiến dữ liệu đầu vào/đầu ra có các phân đoạn không thể phân chia, vì vậy, lượt truyền này sẽ cập nhật các phân đoạn đó thành tiền tố phân đoạn phương diện lớn nhất của phân đoạn ban đầu được phân đoạn đều nhau.