Mục tiêu dài hạn là biến Shardy thành một thành phần hoàn toàn độc lập, có thể hoạt động với bất kỳ phương ngữ MLIR nào. Hiện tại, Shardy trực tiếp phụ thuộc vào StableHLO, nhưng chúng tôi đang tiến hành nâng cấp thông qua nhiều giao diện và khái niệm trừu tượng để giúp Shardy linh hoạt hơn.
Quy tắc phân đoạn
Quy tắc phân đoạn mã hoá cách chúng ta truyền tải thông qua một thao tác. Vì Shardy hiện phụ thuộc vào StableHLO, nên Shardy xác định các quy tắc phân đoạn cho mỗi thao tác stablehlo. Ngoài ra, Shardy cung cấp ShardingRuleOpInterface mà chủ sở hữu phương ngữ có thể sử dụng trong các hoạt động của họ để xác định các quy tắc phân đoạn cho các hoạt động của riêng họ. Miễn là một thao tác triển khai giao diện này, Shardy sẽ có thể truyền tải thông qua giao diện đó.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Thao tác luồng dữ liệu
Một số toán tử, ví dụ: toán tử dựa trên khu vực, yêu cầu một phương pháp khác, trong đó các quy tắc phân đoạn chỉ mô tả mối tương ứng giữa các phương diện trên tất cả toán hạng và kết quả là chưa đủ. Trong những trường hợp này, Shardy xác định một ShardableDataFlowOpInterface để chủ sở hữu phương ngữ có thể mô tả việc truyền tải hoạt động phân đoạn thông qua các thao tác của họ. Giao diện này cung cấp các phương thức để lấy nguồn và đích của mỗi cạnh luồng dữ liệu thông qua chủ sở hữu của chúng, đồng thời lấy và đặt các phân đoạn của chủ sở hữu cạnh.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Xem thêm phần Thao tác luồng dữ liệu để biết thông tin tổng quan về cách chúng tôi xử lý thao tác luồng dữ liệu.
Giao diện chưa được triển khai
Trong tương lai, chúng tôi sẽ thêm nhiều giao diện và đặc điểm hơn để giúp Shardy linh hoạt hơn và không phụ thuộc vào phương ngữ. Chúng tôi liệt kê các trường hợp đó ở bên dưới.
Phân tách hằng số
Hầu hết các chương trình tensor trong MLIR đều có một thực thể của hằng số được sử dụng lại bởi bất kỳ toán tử nào cần giá trị đó. Điều này có ý nghĩa khi hằng số cần thiết là giống nhau. Tuy nhiên, để phân đoạn một chương trình một cách tối ưu, chúng ta muốn cho phép mỗi lần sử dụng một hằng số có phân đoạn riêng và không bị ảnh hưởng bởi cách các toán tử khác sử dụng hằng số đó.
Ví dụ trong hình dưới đây, nếu add được phân đoạn, thì điều này sẽ không ảnh hưởng đến cách phân đoạn divide và subtract (ở các phần khác nhau của phép tính).
Chúng tôi gọi đây là phần phụ thuộc giả: vì các hằng số có giá trị nhỏ, nên không có phần phụ thuộc thực sự nào giữa các toán tử sử dụng cùng một hằng số. Do đó, người dùng có thể quyết định phân đoạn các toán tử hằng số (và toán tử giống hằng số) của họ. Sau đó, mỗi lần sử dụng hằng số đó có thể có một phân đoạn khác nhau có thể tự sao chép bản sao của phép tính phụ hằng số.
Để đạt được điều này, người dùng Shardy cần xác định: – Lệnh truyền your_dialect.constant –> sdy.constant; – Đặc điểm sdy::ConstantLike, chẳng hạn như iota; – Đặc điểm mlir::Elementwise cho các thao tác theo phần tử như add và multiply; – sdy::ConstantFoldable cho các thao tác như slice/broadcast.
Về mặt kỹ thuật, các toán tử này có thể được tính toán tại thời điểm biên dịch, nếu tất cả toán hạng/kết quả của chúng đều là hằng số.
Mức độ ưu tiên của thao tác
Trong GSPMD, các toán tử theo phần tử được truyền trước, theo sau là các toán tử như matmul.
Trong Shardy, chúng tôi muốn cho phép người dùng đặt mức độ ưu tiên cho các thao tác của riêng họ vì chúng tôi không biết trước về các phương ngữ của họ. Do đó, chúng ta sẽ yêu cầu họ truyền một danh sách các thao tác theo thứ tự mà họ muốn Shardy truyền các thao tác đó.
Hình dưới đây cho thấy cách các mức độ ưu tiên được sử dụng trong GSPMD để truyền các thao tác theo thứ tự phù hợp.
Hãy xem bài viết về GSMPD để thảo luận về lý do các mức độ ưu tiên của thao tác là quan trọng.
Không phân biệt phương ngữ
Miễn là bạn triển khai các giao diện, đặc điểm và thẻ trước đó, Shardy sẽ có thể hoạt động cho phương ngữ của bạn. Chúng tôi đang nỗ lực để Shardy trở nên linh hoạt hơn và không phụ thuộc vào phương ngữ. Hãy chú ý theo dõi để nắm bắt thông tin cập nhật về sau.