Shardy, सभी बोलियों के लिए MLIR पर आधारित टेंसर पार्टिशनिंग सिस्टम है. इसे GSPMD और PartIR, दोनों टीमों के साथ मिलकर बनाया गया है. इसमें दोनों सिस्टम के सबसे अच्छे फ़ीचर शामिल हैं. साथ ही, इसमें दोनों टीमों और उपयोगकर्ताओं के अनुभव को भी शामिल किया गया है.
फ़ायदे
- GSPMD के प्रॉपेगेशन को PartIR की इंक्रीमेंटल पार्टिशनिंग के साथ जोड़कर, उपयोगकर्ताओं के लिए ज़्यादा कंट्रोल और अनुमान लगाने की सुविधा.
- शेयर किए गए अनुभव से मिलने वाली नई सुविधाएं.उदाहरण के लिए, आकार बदलने की सुविधा के लिए नई सहायता. यह सुविधा, तब तक ज़्यादा बातचीत जनरेट करती है, जब तक उपयोगकर्ताओं को इसका इस्तेमाल करने का तरीका नहीं पता होता.
- बेहतर इस्तेमाल और डीबग करने की सुविधा, ताकि असली उपयोगकर्ताओं की संख्या बढ़ाई जा सके. उदाहरण के लिए, अक्ष पर आधारित शर्डिंग का इस्तेमाल करके.
- MLIR का इस्तेमाल करने वाला, आसान और ओपन सोर्स कोडबेस. इसमें उपयोगकर्ताओं की मदद करने के लिए, कई तरह के योगदान देने वाले लोग (इंटरनल, बाहरी, और अलग-अलग टाइम ज़ोन में) शामिल हैं.
घटक
- शीयरिंग का तरीका: यह ऐक्सिस पर आधारित शीयरिंग का तरीका है, जो संभावित तौर पर कई मेश में से किसी एक लॉजिकल मेश से जुड़ा होता है. साथ ही, यह डाइमेंशन शीयरिंग और ऐक्सिस को सीमित करने की सुविधा देता है. साथ ही, रीशेप जैसे ऑपरेशन के लिए ऐक्सिस को अलग-अलग करने, इंक्रीमेंटल पार्टिशनिंग के लिए प्राथमिकताएं तय करने वगैरह की सुविधा देता है.
- कंपाइलर एपीआई: कंपाइलर कॉम्पोनेंट का एक सेट, जिसका इस्तेमाल शीयरिंग के प्रॉपेगेशन पर असर डालने के लिए, शीयरिंग के प्रतिनिधित्व के साथ किया जा सकता है.
- इनपुट/आउटपुट के लिए शर्डिंग - मुख्य फ़ंक्शन के इनपुट या आउटपुट में शर्डिंग अटैच करें, ताकि यह पता चल सके कि फ़ंक्शन में दिए जाने पर/फ़ंक्शन से वापस मिलने पर, इनपुट/आउटपुट टेंसर को इस तरह से शर्ड किया जाना चाहिए.
- शीयर करने से जुड़ी शर्त - किसी इंटरमीडिएट टेंसर (उदाहरण के लिए, मैटमल का नतीजा) में शीयर करने की सुविधा जोड़ें, ताकि यह पता चल सके कि उस टेंसर या उसके इस्तेमाल के सबसेट को इस तरह से शीयर किया जाना चाहिए.
- इस तरह से/जैसे सेगमेंट में बांटें - एक आईडी के हिसाब से कई टेंसर को ग्रुप करें, ताकि यह पता चल सके कि उन्हें एक ही तरह से सेगमेंट में बांटा जाना चाहिए.
- मैन्युअल कैलकुलेशन - इसमें एक सब-कैलकुलेशन शामिल होता है, जिसे मैन्युअल रूप से, मेश ऐक्सिस के सबसेट का इस्तेमाल करके बांटा जाता है.इसमें, उन मैन्युअल ऐक्सिस के साथ शर्डिंग, सभी इनपुट और आउटपुट के लिए तय की जाती है.साथ ही, सब-कैलकुलेशन में टेंसर टाइप, उन शर्डिंग के हिसाब से लोकल होते हैं.
- शीयरिंग प्रोपेगेशन: यह एक ऐसा प्रोपेगेशन एल्गोरिदम है जो उपयोगकर्ता की प्राथमिकताओं और शीयरिंग की सीमाओं को, कंपाइलर के लागत-मॉडल और हेयुरिस्टिक्स के साथ जोड़ता है:
- उपयोगकर्ता की तय प्राथमिकताएं, जैसे कि बैच पैरलललिज़्म के बाद ZeRO का इस्तेमाल करना
- ऑपरेशन के आधार पर प्राथमिकताएं, जैसे कि पहले एलिमेंट के हिसाब से ऑपरेशन और फिर मैटमल वगैरह.
- बेहतर तरीके से काम करने वाले हेयुरिस्टिक्स, जैसे कि बैच डाइमेंशन को प्राथमिकता देना.
- SPMD Partitioner: यह एक ऐसा कॉम्पोनेंट है जो प्रोग्राम को SPMD प्रोग्राम में बांटकर, डेटा को एक से दूसरी जगह भेजने/फ़ॉर्मैट करने और प्रोसेस में एक साथ कई कार्रवाइयां करने की ज़रूरत को कम करता है.
- शुरुआती लागू करने के लिए, मौजूदा GSPMD SPMD के partitioner का इस्तेमाल किया जाएगा.
- आने वाले समय में, हमारा मकसद MLIR पर आधारित नया SPMD पार्टीशनर बनाना है.
कोड का डेटा स्टोर करने की जगह
Shardy प्रोजेक्ट पर लगातार काम किया जा रहा है. हम ओपन सोर्स कम्यूनिटी से सुझाव, राय या शिकायत चाहते हैं. Shardy कोड https://github.com/openxla/shardy पर उपलब्ध है