Shardy là một hệ thống phân vùng tensor dựa trên MLIR cho tất cả các phương ngữ. Được xây dựng từ sự cộng tác của cả nhóm GSPMD và nhóm PartIR, công cụ này kết hợp những tính năng tốt nhất của cả hai hệ thống và trải nghiệm chung của cả nhóm và người dùng.
Lợi ích
- Người dùng có nhiều quyền kiểm soát và khả năng dự đoán hơn bằng cách kết hợp việc truyền tải của GSPMD với tính năng phân vùng gia tăng của PartIR.
- Các tính năng mới dựa trên trải nghiệm được chia sẻ, ví dụ: tính năng hỗ trợ mới cho việc định hình lại thường tạo ra nhiều thông tin liên lạc hơn trừ phi người dùng biết cách xử lý các tính năng này.
- Cải thiện khả năng hữu dụng và gỡ lỗi để tăng tốc độ của người dùng cuối, ví dụ: bằng cách sử dụng bản trình bày phân đoạn dựa trên trục.
- Cơ sở mã nguồn mở, đơn giản sử dụng MLIR, với một nhóm cộng tác viên tích cực rộng hơn (nội bộ, bên ngoài và trên nhiều múi giờ) để hỗ trợ người dùng.
Thành phần
- Biểu thị phân đoạn: Biểu thị phân đoạn dựa trên trục được liên kết với một lưới logic cụ thể (trong số nhiều lưới có thể có) và hỗ trợ các trục và phân đoạn kích thước ràng buộc, phân tách trục cho các thao tác như định hình lại, mức độ ưu tiên cho việc phân vùng gia tăng, v.v.
- API trình biên dịch: Một tập hợp các thành phần trình biên dịch có thể được sử dụng cùng với bản trình bày phân đoạn để ảnh hưởng đến việc truyền phân đoạn.
- Phân đoạn đầu vào/đầu ra – đính kèm một phân đoạn vào đầu vào hoặc đầu ra của hàm chính để cho biết đây là cách phân đoạn tensor đầu vào/đầu ra khi được cung cấp cho/trả về từ hàm.
- Quy tắc ràng buộc phân đoạn – đính kèm một phân đoạn vào một tensor trung gian (ví dụ: kết quả của matmul) để cho biết đây là cách phân đoạn tensor đó hoặc một tập hợp con của các cách sử dụng tensor đó.
- Phân đoạn dưới dạng/giống như – nhóm nhiều tensor theo mã nhận dạng để cho biết rằng các tensor đó phải được phân đoạn theo cùng một cách.
- Tính toán thủ công – Bao gồm một phép tính phụ được phân vùng theo cách thủ công bằng cách sử dụng một tập hợp con các trục lưới, trong đó các phân đoạn dọc theo các trục thủ công đó được chỉ định cho tất cả dữ liệu đầu vào và đầu ra, và bên trong phép tính phụ, các loại tensor là cục bộ so với các phân đoạn đó.
- Truyền tải phân đoạn: Thuật toán truyền tải kết hợp các mức độ ưu tiên của người dùng và các quy tắc ràng buộc phân đoạn với mô hình chi phí trình biên dịch và các phương pháp phỏng đoán:
- Mức độ ưu tiên do người dùng xác định, ví dụ: thực hiện song song theo lô rồi đến ZeRO
- Mức độ ưu tiên dựa trên toán tử, ví dụ: toán tử theo phần tử trước, sau đó là toán tử nhân ma trận, v.v.
- Các phương pháp phỏng đoán chi tiết hơn, ví dụ: ưu tiên phương diện theo lô.
- Trình phân vùng SPMD: Một thành phần giúp giảm các quyết định truyền tải phân đoạn bằng cách phân vùng chương trình thành một chương trình SPMD, thêm thao tác di chuyển/định dạng dữ liệu cần thiết và các thao tác tập thể trong quá trình này.
- Trong ngắn hạn, quá trình triển khai ban đầu sẽ sử dụng trình phân vùng SPMD của GSPMD hiện tại.
- Về lâu dài, chúng tôi dự định tạo một trình phân vùng SPMD mới dựa trên MLIR.
Kho lưu trữ mã
Dự án Shardy đang trong quá trình phát triển và chúng tôi muốn nhận được ý kiến phản hồi của cộng đồng nguồn mở. Bạn có thể xem mã Shardy tại https://github.com/openxla/shardy