लंबे समय तक का लक्ष्य है कि Shardy को पूरी तरह से स्टैंडअलोन कॉम्पोनेंट बनाया जाए, जो किसी भी MLIR डायलेक्ट के साथ काम कर सके. फ़िलहाल, Shardy सीधे तौर पर StableHLO पर निर्भर है. हालांकि, हम Shardy को ज़्यादा सुविधाजनक बनाने के लिए, अलग-अलग एब्स्ट्रैक्शन और इंटरफ़ेस की मदद से, इस निर्भरता को कम करने की कोशिश कर रहे हैं.
डेटा का बंटवारा करने के नियम
शर्डिंग नियम से यह तय होता है कि हम किसी ऑपरेशन के ज़रिए डेटा को कैसे प्रोपेगेट करते हैं. Shardy अब StableHLO पर निर्भर करता है. इसलिए, यह हर stablehlo ऑपरेशन के लिए, sharding के नियम तय करता है. इसके अलावा, Shardy ShardingRuleOpInterface
उपलब्ध कराता है. इसका इस्तेमाल, बोली के मालिक अपने ऑपरेशन में कर सकते हैं, ताकि वे अपने ऑपरेशन के लिए, sharding के नियम तय कर सकें. जब तक कोई कार्रवाई इस इंटरफ़ेस को लागू करती है, तब तक Shardy इस इंटरफ़ेस के ज़रिए प्रोपेगेट कर पाएगा.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
डेटा फ़्लो ऑपरेशंस
कुछ ऑपरेशन, जैसे कि क्षेत्र के हिसाब से ऑपरेशन के लिए, अलग तरीके की ज़रूरत होती है. ऐसे ऑपरेशन के लिए, सिर्फ़ सभी ऑपरेंड और नतीजों के डाइमेंशन के बीच के संबंध के बारे में बताने वाले, शार्डिंग नियम काफ़ी नहीं होते. ऐसे मामलों में, Shardy एक ShardableDataFlowOpInterface
तय करता है, ताकि बोली के मालिक अपने ऑपरेशंस के ज़रिए, sharding के प्रॉपेगेशन के बारे में बता सकें. यह इंटरफ़ेस, हर डेटाफ़्लो एज के सोर्स और टारगेट को उनके मालिक के ज़रिए पाने के तरीके उपलब्ध कराता है. साथ ही, एज के मालिकों की शार्डिंग को पाने और सेट करने की सुविधा भी देता है.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
डेटा फ़्लो ऑपरेशंस को मैनेज करने के तरीके की खास जानकारी पाने के लिए, डेटा फ़्लो ऑपरेशंस भी देखें.
इंटरफ़ेस अभी तक लागू नहीं किए गए हैं
आने वाले समय में, Shardy को ज़्यादा बेहतर और आसान बनाने के लिए, इसमें और इंटरफ़ेस और ट्रेट जोड़े जाएंगे. साथ ही, यह बोली के हिसाब से भी काम करेगा. हमने इनकी सूची यहां दी है.
लगातार अलग-अलग ग्रुप में बांटना
MLIR में ज़्यादातर टेंसर प्रोग्राम में, किसी कॉन्स्टेंट का एक इंस्टेंस होता है. जिस ऑपरेशन में इस वैल्यू की ज़रूरत होती है वह इसका फिर से इस्तेमाल करता है. यह तब सही होता है, जब एक ही कॉन्स्टेंट की ज़रूरत हो. हालांकि, किसी प्रोग्राम को ऑप्टिमाइज़ तरीके से शर्ड करने के लिए, हम चाहते हैं कि किसी कॉन्स्टेंट के हर इस्तेमाल के लिए अलग-अलग शर्डिंग की सुविधा हो. साथ ही, इस बात से कोई फ़र्क़ न पड़े कि दूसरे ऑपरेशन उस कॉन्स्टेंट का इस्तेमाल कैसे करते हैं.
उदाहरण के लिए, नीचे दिए गए चित्र में, अगर add
को अलग-अलग हिस्सों में बांटा गया है, तो इसका असर इस बात पर नहीं पड़ना चाहिए कि divide
और subtract
(कंप्यूटेशन के अलग-अलग हिस्सों में) को कैसे अलग-अलग हिस्सों में बांटा गया है.
हम इसे गलत डिपेंडेंसी कहते हैं: कॉन्स्टेंट कम जगह लेते हैं. इसलिए, एक ही कॉन्स्टेंट का इस्तेमाल करने वाले ऑपरेशन के बीच कोई असल डिपेंडेंसी नहीं होती. इसलिए, उपयोगकर्ता अपने कॉन्स्टेंट (और कॉन्स्टेंट जैसे) ऑपरेशन को अलग-अलग हिस्सों में बांटने का फ़ैसला कर सकते हैं. इसके बाद, उस कॉन्स्टेंट के हर इस्तेमाल के लिए अलग-अलग शर्डिंग की जा सकती है. यह शर्डिंग, कॉन्स्टेंट के सब-कंप्यूटेशन की अपनी कॉपी में अलग से प्रोपेगेट की जा सकती है.
ऐसा करने के लिए, Shardy के उपयोगकर्ताओं को ये तय करने होंगे: - your_dialect.constant
->
sdy.constant
पास; - sdy::ConstantLike
ट्रीट, जैसे कि
iota; - add
और
multiply
जैसे एलिमेंट के हिसाब से ऑपरेशन के लिए mlir::Elementwise
ट्रीट; - slice/broadcast जैसे ऑपरेशन के लिए
sdy::ConstantFoldable
.
अगर इन ऑपरेशन के सभी ऑपरेंड/नतीजे कॉन्स्टेंट हैं, तो तकनीकी तौर पर इनकी गिनती कंपाइल करने के समय की जा सकती है.
ऑपरेशन की प्राथमिकताएं
GSPMD में, एलिमेंट के हिसाब से ऑपरेशन पहले प्रॉगेट किए जाते हैं. इसके बाद, matmul
जैसे ऑपरेशन प्रॉगेट किए जाते हैं.
Shardy में, हम उपयोगकर्ताओं को अपनी प्राथमिकताएं सेट करने की अनुमति देना चाहते हैं, क्योंकि हमें उनकी बोलियों के बारे में पहले से जानकारी नहीं है. इसलिए, हम उनसे ऑपरेशन की सूची को उसी क्रम में पास करने के लिए कहेंगे जिस क्रम में Shardy को उन्हें प्रोपेगेट करना है.
नीचे दिए गए इलस्ट्रेशन में दिखाया गया है कि ऑपरेशन को सही क्रम में प्रोपेगेट करने के लिए, GSPMD में प्राथमिकताओं का इस्तेमाल कैसे किया जाता है.
प्राथमिकताएं तय करना क्यों ज़रूरी है, इस बारे में जानने के लिए GSPMD पेपर देखें.
बोली के हिसाब से अलग-अलग नहीं होना
जब तक आपने पिछले इंटरफ़ेस, ट्रेट, और पास लागू किए हैं, तब तक Shardy आपकी बोली के हिसाब से काम करेगा. हम Shardy को ज़्यादा सुविधाजनक और बोली के हिसाब से काम करने वाला बनाने पर काम कर रहे हैं. इस बारे में और अपडेट पाने के लिए, हमारे साथ बने रहें.