Compiler API

बैकग्राउंड

हम मानते हैं कि पाठकों को शर्डिंग के बारे में जानकारी है. इससे पता चलता है कि किसी टेंसर को Shardy में कैसे शर्ड किया जा सकता है. इस दस्तावेज़ में बताया गया है कि किसी प्रोग्राम में, शीयरिंग रेप्रज़ेंटेशन का इस्तेमाल कैसे किया जा सकता है. उदाहरण के लिए, प्रोग्राम के किसी खास टेंसर में शीयरिंग अटैच करना.

शेर्डिंग प्रॉपेगेशन, किसी प्रोग्राम में हर टेंसर के लिए शेर्डिंग तय करने की प्रोसेस है. इसमें, टेंसर के सबसेट के लिए शेर्डिंग की सीमाएं तय की जाती हैं. Shardy के compiler API की मदद से, डेटा को अलग-अलग हिस्सों में बांटने की प्रोसेस को कंट्रोल किया जा सकता है. इसके अलावा, यह उपयोगकर्ताओं को अपने प्रोग्राम में, मैन्युअल तरीके से अलग-अलग हिस्सों में बांटकर किए गए कैलकुलेशन डालने की अनुमति देता है.

मकसद

इस दस्तावेज़ में, Shardy में ऐसे एपीआई कॉम्पोनेंट के डिज़ाइन के बारे में बताया गया है. साथ ही, उनके व्यवहार और इनवैरिएंट के बारे में भी बताया गया है. ध्यान दें कि इस एपीआई का इस्तेमाल, शीयरिंग के प्रॉपेगेशन को कंट्रोल करने के लिए किया जाता है. हालांकि, इस दस्तावेज़ में प्रॉपेगेशन के व्यवहार या इसके डिज़ाइन के बारे में नहीं बताया गया है.

खास जानकारी

• इनपुट/आउटपुट के लिए शीयरिंग - मुख्य फ़ंक्शन के इनपुट या आउटपुट में शीयरिंग अटैच करें, ताकि यह पता चल सके कि फ़ंक्शन में दिए जाने पर/फ़ंक्शन से वापस मिलने पर, इनपुट/आउटपुट टेंसर को इस तरह से शीयर किया जाना चाहिए.

• शर्ड करने से जुड़ी शर्त - किसी इंटरमीडिएट टेंसर (जैसे, matmul का नतीजा) में शर्डिंग अटैच करें, ताकि यह पता चल सके कि उस टेंसर या उसके इस्तेमाल के सबसेट को कैसे शर्ड किया जाना चाहिए.

• शर्डिंग ग्रुप - एक आईडी के हिसाब से कई टेंसर को ग्रुप करें, ताकि यह पता चल सके कि उन्हें एक ही तरह से शर्ड किया जाना चाहिए.

• मैन्युअल कैलकुलेशन - इसमें एक सब-कैलकुलेशन शामिल होती है, जिसे मेश ऐक्सिस के सबसेट का इस्तेमाल करके मैन्युअल रूप से बांटा जाता है.इसमें, उन मैन्युअल ऐक्सिस के साथ shardings को सभी इनपुट और आउटपुट के लिए तय किया जाता है.साथ ही, सब-कैलकुलेशन में, उन shardings के हिसाब से टेंसर टाइप स्थानीय होते हैं.

ज़्यादा जानकारी वाला डिज़ाइन

इनपुट/आउटपुट के लिए अलग-अलग हिस्से

उपयोगकर्ताओं को मुख्य फ़ंक्शन के इनपुट और आउटपुट के लिए, एक शीयरिंग तय करने की अनुमति मिलती है.

MLIR में, एट्रिब्यूट को फ़ंक्शन के आर्ग्युमेंट और नतीजों से अटैच किया जा सकता है. इसलिए, उपयोगकर्ता इस तरह से फ़ंक्शन में, शर्डिंग एट्रिब्यूट अटैच कर सकते हैं.

उदाहरण के लिए:

@mesh_xy = <["x"=2, "y"=2]>

// The 1st input has a sharding specified, but the 2nd input doesn't.
// The output has a sharding specified.
func @main(%arg0: tensor<8x8xf32>
            {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"}, {}]>},
            %arg1: tensor<8x16xf32>)
    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_xy, [{}, {"y"}]>}) {
  ...
}

शार्डिंग की सीमा

इसकी मदद से, उपयोगकर्ता अपने प्रोग्राम में किसी इंटरमीडिएट टेंसर में, शीयरिंग जोड़ सकते हैं. इससे, पार्टीशनर को यह पता चलता है कि टेंसर या उसके इस्तेमाल के सबसेट को इस तरह से शीयर किया जाना चाहिए.

यह एक MLIR ऑपरेशन है, जो टेंसर को इनपुट के तौर पर लेता है. साथ ही, इसमें एक sharding एट्रिब्यूट जुड़ा होता है. ऑपरेशन इनमें से कोई एक हो सकता है:

• जिनका इस्तेमाल नहीं किया गया है (डैंगल) - इसका मतलब है कि अटैच की गई sharding के हिसाब से, टेन्सर को भी sharding किया जाना चाहिए.
• इस्तेमाल किया गया है - इसका मतलब है कि अटैच की गई sharding से पता चलता है कि sharding constraint op के इस्तेमाल को कैसे sharding किया जाना चाहिए. हालांकि, इनपुट टेंसर के अन्य इस्तेमाल के लिए अलग sharding हो सकती है. अगर इनपुट टेंसर का कोई और इस्तेमाल नहीं किया गया है, तो इसका व्यवहार, 'इस्तेमाल नहीं किया गया' केस के जैसा ही होगा. प्रॉपेगेशन, टेन्सर के लिए अपने-आप sharding तय करेगा और ज़रूरत पड़ने पर उसे फिर से sharding करेगा.

इसमें डाइमेंशन की ओपन शर्डिंग हो सकती है. इसका मतलब है कि ऑपरेंड को उपलब्ध ऐक्सिस के साथ और भी शर्ड किया जा सकता है.

@mesh_xy = <["x"=2, "y"=2]>

%0 = ... : tensor<8x8xf32>
%1 = sdy.sharding_constraint %0 <@mesh_xy, [{"x"}, {?}]> : tensor<8x8xf32>

शीयर करने वाला ग्रुप

जिन मामलों में दो या उससे ज़्यादा टेंसर के बीच कोई डेटा डिपेंडेंसी या कोई स्ट्रॉन्ग डेटा डिपेंडेंसी नहीं होती है, लेकिन उपयोगकर्ताओं को पता होता है कि उन टेंसर को एक जैसे या मिलते-जुलते तरीके से पार्टिशन किया जाना चाहिए, तो Shardy API इस संबंध को बताने का एक तरीका उपलब्ध कराता है. इससे उपयोगकर्ताओं को साफ़ तौर पर यह बताने की सुविधा मिलती है कि टेंसर को एक-दूसरे के तौर पर बांटा जाना चाहिए.

ऐसा करने के लिए, हमने शर्ड ग्रुप की सुविधा शुरू की है. इसमें हर ग्रुप में, एक ही शर्ड ग्रुप आईडी से जुड़े कई निर्देश होते हैं. एक ही ग्रुप में, एक जैसे हिस्से बनाने के लिए, ग्रुप को अलग-अलग हिस्सों में बांटने की सुविधा का इस्तेमाल किया जाता है.

उदाहरण के लिए, नीचे दिखाए गए काल्पनिक उपयोगकर्ता प्रोग्राम में, हम प्रोग्राम के आउटपुट को ठीक वैसे ही शेयर करना चाहते हैं जैसे प्रोग्राम के इनपुट को. हालांकि, इन दोनों के बीच डेटा की कोई डिपेंडेंसी नहीं है.

अगर हम इस प्रोग्राम को चलाते हैं, तो शीयरिंग प्रॉपेगेशन, टेन्सर %1 और %2 की शीयरिंग के बारे में अनुमान नहीं लगा पाएगा. साथ ही, इनका दोहराव हो जाएगा. हालांकि, shard_group एट्रिब्यूट को अटैच करके, हम यह बताते हैं कि इनपुट %0 और आउटपुट %2 एक ही shard_group में हैं. इससे, @mesh_xy, [{"x"},{"y"}]> को इनपुट %0 से आउटपुट %2 में और फिर बाकी ग्राफ़ में भेजा जा सकता है. यहां ग्राफ़ को लगातार %1 ब्रॉडकास्ट किया जाता है. हम sdy.sharding_group ऑपरेशन की मदद से, किसी ग्रुप को वैल्यू असाइन कर सकते हैं.

@mesh_xy = <["x"=2, "y"=2]>

module @"jit_zeros_like" {
  func.func @main(%arg0: tensor<8x2xi64> {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"},{"y"}]>} }) -> (tensor<8x2xi64>) {
    %0 = sdy.sharding_group %arg0, id=0 : tensor<8x2xi64>
    %1 = stablehlo.constant dense<0> : tensor<8x2xi64>
    %2 = sdy.sharding_group %1, id=0 : tensor<8x2xi64>
    return %2 : tensor<8x2xi64>
  }
}

ऊपर दिए गए इस आसान उदाहरण में, हमने आउटपुट के लिए वही sharding तय की है जो इनपुट के लिए तय की थी. इससे वही नतीजा मिलेगा, क्योंकि हमें पहले से पता है कि इनपुट के लिए कौनसी sharding असाइन करनी है. हालांकि, ज़्यादा काम के उदाहरणों में, हम sharding का इस्तेमाल करते हैं, ताकि कई टेंसर की sharding को सिंक किया जा सके. इसके लिए, यह ज़रूरी नहीं है कि हम उनमें से किसी के लिए भी sharding तय करें. Shardy बाकी काम करेगा और उन्हें असाइन करने के लिए सबसे अच्छी sharding ढूंढेगा.

मैन्युअल तरीके से कैलकुलेट करना

हो सकता है कि उपयोगकर्ता यह कंट्रोल करना चाहें कि उनके कैलकुलेशन के हिस्सों को कैसे बांटा जाए और किन कलेटिव का इस्तेमाल किया जाए. उदाहरण के लिए, कुछ उपयोगकर्ता कंपाइलर को देर करने के बजाय, फ़्रंटएंड एपीआई से मैन्युअल तरीके से, एक साथ कई मैटमल लागू करना चाहते हैं. हम मैन्युअल कैलकुलेशन एपीआई उपलब्ध कराते हैं, जिसकी मदद से वे ऐसा कर सकते हैं.

यह मैन्युअल सब-कंप्यूटेशन के लिए, एक क्षेत्र वाला MLIR ऑपरेशन है. उपयोगकर्ता, मेश ऐक्सिस के सबसेट (संभवतः सभी) का इस्तेमाल करके, इस सब-कंप्यूटेशन के लिए इनपुट/आउटपुट की शर्डिंग तय करेंगे. सब-कंप्यूटेशन, तय किए गए मेश ऐक्सिस (जिन्हें मैन्युअल ऐक्सिस भी कहा जाता है) के हिसाब से लोकल/मैन्युअल होगा. साथ ही, तय नहीं किए गए ऐक्सिस (जिन्हें फ़्री ऐक्सिस भी कहा जाता है) के हिसाब से ग्लोबल/बिना बंटवारे वाला होगा. प्रॉपेगेशन के दौरान, सब-कंप्यूटेशन को फ़्री ऐक्सिस के साथ-साथ, उसी तरह से शर्ड किया जा सकता है जिस तरह इस ऑपरेशन से बाहर के कंप्यूटेशन को किया जा सकता है.

उदाहरण के लिए:

@mesh_name = <["data"=2, "model"=2]>

%0 = ... : tensor<16x32xf32>
%1 = sdy.manual_computation(%0)
    in_shardings=[<@mesh_name, [{"data"}, {"model",?}]>]
    out_shardings=[<@mesh_name, [{"data"}, {?}]>]
    manual_axes={"data"}
    (%arg1: tensor<8x32xf32>) {
  // body
  return %42 : tensor<8x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

इनवैरिएंट

1. सभी in_shardings, out_shardings, और manual_axes एक ही नेटवर्क से जुड़े होने चाहिए. manual_axes को मेश के हिसाब से क्रम में लगाया गया है.

2. manual_axes का इस्तेमाल, इन/आउट सभी शार्डिंग में साफ़ तौर पर किया जाना चाहिए. इसका मतलब है कि हर शार्डिंग के लिए, सभी मैन्युअल ऐक्सिस को या तो किसी डाइमेंशन को शार्ड करना होगा या साफ़ तौर पर डुप्लीकेट किया जाना होगा.

3. अगर किसी इन/आउट sharding में कोई फ़्री ऐक्सिस (manual_axes में मौजूद कोई भी मेश ऐक्सिस) मौजूद है, तो वह उसी डाइमेंशन sharding में मौजूद किसी भी मैन्युअल ऐक्सिस से छोटा होना चाहिए. ऊपर दिए गए उदाहरण में, डाइमेंशन sharding {"model", "data"} अमान्य होगी.

4. कैलकुलेशन का क्षेत्र/बॉडी, स्थानीय कैलकुलेशन होता है. जैसे, उपयोगकर्ता के तय किए गए कलेक्शन शामिल करना. यह मैन्युअल ऐक्सिस के साथ इन/आउट sharding के हिसाब से स्थानीय होना चाहिए (ऊपर दिया गया नोट देखें).

मैन्युअल तरीके से किए गए कैलकुलेशन को नेस्ट करना

एक-दूसरे के अंदर कई मैन्युअल कैलकुलेशन नेस्ट किए जा सकते हैं. हालांकि, इसके लिए ज़रूरी है कि हर कैलकुलेशन, मैन्युअल ऐक्सिस के अपने यूनीक सेट पर काम करे.