-sdy-close-shardings
Đóng các phân đoạn tensor và thả các trục được sao chép.
-sdy-constant-merger
Hợp nhất các hằng số giống hệt nhau với các phân đoạn trùng khớp.
Thực hiện CSE nhẹ trên các hằng số có phân đoạn giống hệt nhau.
Quy trình nhập phân tách và sao chép các hằng số để việc phân đoạn không được truyền giữa các cách sử dụng khác nhau của một phép tính phụ hằng số. Nếu các hằng số có cùng phân đoạn sau khi truyền, thì lượt truyền này sẽ hợp nhất các hằng số đó để tiết kiệm thời gian biên dịch.
-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 yêu cầu bất kỳ hoạt động giao tiếp phân đoạn lại nào (lưu ý rằng toán tử vẫn có thể yêu cầu giao tiếp 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 truyền này, các thao tác sẽ không 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ụ:
Đầu vào:
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>
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, lhs
và rhs
đều được phân đoạn theo trục "x" trên các kích thước không co giãn, không tương thích. Lệnh truyền này sẽ chèn một phân đoạn lại rõ ràng trên rhs
trước toán tử dấu chấm để toán tử dấu chấm có các phân đoạn tương thích.
-sdy-remove-sharding-groups
Xoá ShardingGroupOps sau khi truyền.
-sdy-reshard-to-collectives
Chuyển đổi ReshardOp thành nhiều toán tử tập hợp Shardy.
So khớp các thao tác phân đoạn lại và viết lại các thao tác đó thành nhiều thao tác tập hợp Shardy. Sau lần truyền này, không còn thao tác phân đoạn lại nào trong mô-đun. Lệnh truyền này giả định rằng các phân đoạn rõ ràng đã được chèn (sdy-insert-explicit-reshards
).
Ví dụ:
Đầu vào:
mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>
Kết quả:
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>
Trong ví dụ trên, tensor %0 : tensor<16x2xf32>
được phân đoạn thành \[{"x", "y", "z"}, {}\]
. Sau đó, có một thao tác reshard
phân đoạn lại dưới dạng \[{"x"}, {}\]
. Trên trục đầu tiên, vì hậu tố {"y", "z"}
bị xoá sau khi phân đoạn lại, nên chúng ta suy luận rằng chúng ta đã thu thập tất cả {"y", "z"}
. Phương diện thứ hai không thay đổi.
-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.
Tùy chọn
-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-temp-explicit-reshards-for-optimizations
Chèn các phân đoạn lại rõ ràng cho các hoạt động tối ưu hoá cụ thể.
Thẻ và vé này là giải pháp tạm thời cho đến khi chúng ta có thể bật thẻ và vé sdy-insert-explicit-reshards
theo mặc định.
Điều này cho phép chúng ta chèn các phân đoạn lại rõ ràng trên các thao tác cụ thể để tối ưu hoá.
-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.