Sharding Agnostik Dialek

Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.

Sasaran jangka panjangnya adalah menjadikan Shardy sebagai komponen yang sepenuhnya mandiri, yang dapat berfungsi dengan dialek MLIR apa pun. Saat ini, Shardy secara langsung bergantung pada StableHLO, tetapi kami membuat progres untuk mengangkatnya melalui berbagai abstraksi dan antarmuka untuk membuat Shardy lebih fleksibel.

Aturan Sharding

Aturan sharding mengenkode cara kami melakukan propagasi melalui operasi. Karena Shardy sekarang bergantung pada StableHLO, Shardy menentukan aturan sharding untuk setiap operasi stablehlo. Selain itu, Shardy menyediakan ShardingRuleOpInterface yang dapat digunakan oleh pemilik dialek dalam operasi mereka untuk menentukan aturan sharding untuk operasi mereka sendiri. Selama operasi mengimplementasikan antarmuka ini, Shardy akan dapat menyebar melalui antarmuka tersebut.

def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
  let methods = [
    InterfaceMethod<
      /*desc=*/[{
        Returns the sharding rule of the op.
      }],
      /*retType=*/"mlir::sdy::OpShardingRuleAttr",
      /*methodName=*/"getShardingRule"
    >,
  ];
}

Operasi aliran data

Beberapa operasi, misalnya, operasi berbasis region, memerlukan pendekatan yang berbeda jika aturan sharding, yang hanya menjelaskan korespondensi antara dimensi di semua operand dan hasil, tidak memadai. Dalam kasus ini, Shardy menentukan ShardableDataFlowOpInterface sehingga pemilik dialek dapat menjelaskan penyebaran sharding melalui operasi mereka. Antarmuka ini menyediakan metode untuk mendapatkan sumber dan target setiap tepi alur data melalui pemiliknya, dan juga mendapatkan serta menetapkan sharding pemilik tepi.

def ShardableDataFlowOpInterface :
    OpInterface<"ShardableDataFlowOpInterface"> {
  (get|set)BlockArgumentEdgeOwnerShardings;
  (get|set)OpResultEdgeOwnerShardings;
  getBlockArgumentEdgeOwners;
  getOpResultEdgeOwners;
  getEdgeSources;
  // ...
}

Lihat juga Operasi aliran data untuk mengetahui ringkasan umum tentang cara kami menangani operasi aliran data.

Antarmuka belum diterapkan

Di masa mendatang, lebih banyak antarmuka dan karakteristik akan ditambahkan untuk membuat Shardy lebih fleksibel dan tidak bergantung pada dialek. Kami mencantumkannya di bawah ini.

Pemisahan konstan

Sebagian besar program tensor di MLIR memiliki satu instance konstanta yang digunakan kembali oleh op apa pun yang memerlukan nilai tersebut. Hal ini masuk akal jika konstanta yang diperlukan sama. Namun, untuk sharding program yang optimal, sebaiknya izinkan setiap penggunaan konstanta memiliki sharding-nya sendiri, dan tidak terpengaruh oleh cara ops lain menggunakan konstanta tersebut.

Misalnya pada gambar di bawah, jika add di-shard, hal ini tidak akan memengaruhi cara divide dan subtract (di bagian komputasi yang berbeda) di-shard.

Kita menyebutnya dependensi palsu: karena konstanta murah, tidak ada dependensi nyata antara operasi yang menggunakan konstanta yang sama. Dengan demikian, pengguna dapat memutuskan sharding operasi konstan (dan mirip konstan) mereka. Setiap penggunaan konstanta tersebut kemudian dapat memiliki sharding yang berbeda yang dapat di-propagasi secara terpisah ke salinan sub-komputasi konstanta-nya sendiri.

Untuk mencapai hal ini, pengguna Shardy perlu menentukan: - your_dialect.constant -> sdy.constant pass; - Sifat sdy::ConstantLike, seperti iota; - Sifat mlir::Elementwise untuk operasi per elemen seperti add dan multiply; - sdy::ConstantFoldable untuk operasi seperti slice/broadcast. Secara teknis, operasi ini dapat dihitung pada waktu kompilasi, jika semua operand/hasilnya adalah konstanta.

Prioritas operasi

Di GSPMD, operasi per elemen disebarkan terlebih dahulu, diikuti dengan operasi seperti matmul. Di Shardy, kami ingin mengizinkan pengguna menetapkan prioritas operasi mereka sendiri karena kami tidak mengetahui dialek mereka a priori. Oleh karena itu, kita akan meminta mereka untuk meneruskan daftar ops dalam urutan yang mereka inginkan agar Shardy menyebarkannya.

Gambar di bawah menunjukkan cara prioritas digunakan di GSPMD untuk menyebarkan operasi dalam urutan yang benar.

Lihat makalah GSMPD untuk mengetahui diskusi tentang mengapa prioritas operasi penting.

Bersifat agnostik dialek

Selama Anda menerapkan antarmuka, karakteristik, dan kartu sebelumnya, Shardy akan dapat berfungsi untuk dialek Anda. Kami berupaya membuat Shardy lebih fleksibel dan tidak bergantung pada dialek, jadi nantikan kabar terbarunya.