Shardy adalah sistem partisi tensor berbasis MLIR untuk semua dialek. Dibuat dari kolaborasi tim GSPMD dan PartIR, fitur ini menggabungkan yang terbaik dari kedua sistem, dan pengalaman bersama dari tim dan pengguna.
Manfaat
- Lebih banyak kontrol dan prediktabilitas untuk pengguna dengan menggabungkan propagasi GSPMD dengan partisi inkremental PartIR.
- Fitur baru yang didorong oleh pengalaman bersama, misalnya dukungan baru untuk pembentukan ulang yang terkenal menghasilkan komunikasi tambahan kecuali jika pengguna tahu cara mengatasinya.
- Kegunaan dan kemampuan debug yang lebih baik untuk meningkatkan kecepatan pengguna akhir, misalnya dengan menggunakan representasi sharding berbasis sumbu.
- Codebase open source sederhana yang menggunakan MLIR, dengan kumpulan kontributor aktif yang lebih luas (internal, eksternal, dan di berbagai zona waktu) untuk mendukung pengguna.
Komponen
- Representasi Sharding: Representasi sharding berbasis sumbu yang terikat ke mesh logika tertentu (dari kemungkinan beberapa mesh), dan mendukung sharding dimensi dan sumbu yang membatasi, memisahkan sumbu untuk operasi seperti membentuk ulang, prioritas untuk partisi inkremental, dan lainnya.
- Compiler API: Kumpulan komponen compiler yang dapat digunakan bersama
representasi sharding untuk memengaruhi penyebaran sharding.
- Sharding input/output - lampirkan sharding ke input atau output fungsi utama, untuk menunjukkan bahwa ini adalah cara tensor input/output harus di-sharding saat diberikan ke/ditampilkan dari fungsi.
- Batasan Sharding - lampirkan sharding ke tensor perantara (misalnya, hasil matmul) untuk menunjukkan bahwa ini adalah cara tensor tersebut, atau subkumpulan penggunaannya, harus di-sharding.
- Shard As/Like - mengelompokkan beberapa tensor menurut ID untuk menunjukkan bahwa tensor tersebut harus di-shard dengan cara yang sama.
- Komputasi Manual - Meliputi sub-komputasi yang dipartisi secara manual menggunakan subset sumbu mesh, dengan sharding di sepanjang sumbu manual tersebut ditentukan untuk semua input dan output, dan di dalam sub-komputasi, jenis tensor bersifat lokal sehubungan dengan sharding tersebut.
- Sharding Propagation: Algoritma propagasi yang menggabungkan prioritas pengguna
dan batasan sharding, dengan model biaya dan heuristik compiler:
- Prioritas yang ditentukan pengguna, misalnya, lakukan paralelisme batch, lalu ZeRO
- Prioritas berbasis operasi, misalnya operasi per elemen terlebih dahulu, lalu matmul, dll.
- Heuristik yang lebih terperinci, misalnya, lebih memilih dimensi batch.
- SPMD Partitioner: Komponen yang menurunkan keputusan penyebaran sharding dengan
mempartisi program menjadi program SPMD, menambahkan pemindahan/pemformatan data
yang diperlukan dan operasi kolektif dalam prosesnya.
- Dalam jangka pendek, implementasi awal akan menggunakan partisi SPMD GSPMD saat ini.
- Dalam jangka panjang, kami berencana membuat partisi SPMD baru berbasis MLIR.
Repositori kode
Project Shardy sedang dalam pengembangan aktif, dan kami mencari masukan dari komunitas open source. Kode Shardy tersedia di https://github.com/openxla/shardy