🤖 machine learning

Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism

यह शोध पत्र $\alpha$ -CROWN सत्यापन ढांचे (verification framework) में टेंसर पैरेललिज्म (Tensor Parallelism) और फुली शार्डेड डेटा पैरेललिज्म (Fully Sharded Data Parallelism) को अनुकूलित करता है ताकि GPU मेमोरी के उपयोग को काफी कम किया जा सके, जिससे CIFAR-100 पर ResNet-large जैसे बड़े पैमाने के न्यूरल नेटवर्क का औपचारिक सत्यापन संभव हो सके जो पहले मेमोरी की सीमाओं के कारण असंभव था।

मूल लेखक: Sergei Vorobyov, Eugene Ilyushin

प्रकाशित 2026-06-09

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Sergei Vorobyov, Eugene Ilyushin

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप यह साबित करने की कोशिश कर रहे हैं कि एक सेल्फ-ड्राइविंग कार कभी दुर्घटनाग्रस्त नहीं होगी, चाहे मौसम कैसा भी हो या कोई पैदल यात्री अचानक सामने आ जाए। आप कार का दस लाख बार परीक्षण नहीं कर सकते; आपको हर संभव स्थिति में कार के सुरक्षित होने का एक गणितीय "प्रमाण" (proof) चाहिए। इसे फॉर्मल न्यूरल नेटवर्क वेरिफिकेशन (Formal Neural Network Verification) कहा जाता है।

यह काम कंप्यूटर मेमोरी पर बहुत भारी होता है। यह एक विशाल पहेली को हल करने जैसा है, लेकिन इसके सभी टुकड़े (डेटा और नियम) एक ही छोटी मेज (एक सिंगल ग्राफिक्स कार्ड) पर फिट होने चाहिए। यदि पहेली बहुत बड़ी है, तो मेज भर जाएगी और प्रमाण विफल हो जाएगा।

यह शोध पत्र इस पहेली को हल करने के दो नए तरीके पेश करता है, जो आज के विशाल AI मॉडल को प्रशिक्षित करने के तरीकों से प्रेरित हैं।

यहाँ उनके दो मुख्य समाधानों का विवरण दिया गया है, जिन्हें सरल उपमाओं के साथ समझाया गया है:

1. "पहेली को विभाजित करने" वाला दृष्टिकोण (टेंसर पैरेललिज्म - Tensor Parallelism)

विचार: कल्पना कीजिए कि आपके पास एक विशाल जिग्सॉ पहेली है। एक व्यक्ति द्वारा पूरी पहेली को पकड़ने के बजाय, आप पहेली को दो हिस्सों में काट देते हैं। व्यक्ति A आधा हिस्सा पकड़ता है, और व्यक्ति B दूसरा आधा हिस्सा। दोनों अपने-अपने टुकड़ों पर काम करते हैं और परिणाम एक-दूसरे को बताते हैं।

यह कैसे काम करता है: शोधकर्ता "वेट्स" (पहेली के टुकड़े) और "नियमों" (गणित) को दो GPUs में विभाजित करते हैं।
अच्छी खबर: यह प्रत्येक कंप्यूटर पर आवश्यक मेमोरी को लगभग आधा (लगभग 2 गुना कमी) कर देता है। यह छोटे या कम गहरे (shallow) पहेलियों के लिए बहुत कुशल है।
चुनौती: जब पहेली गहरी (कई लेयर्स वाली) होती है, तो दोनों व्यक्तियों को पूरे चित्र को देखे बिना अपने हिस्सों के बीच के संबंध का अनुमान लगाना पड़ता है। समय बचाने के लिए, वे बीच के हिस्सों के लिए एक "त्वरित और अनुमानित" विधि (जिसे IBP कहा जाता है) का उपयोग करते हैं।
परिणाम: अंतिम प्रमाण अभी भी सुरक्षित रहता है (यह नहीं कहेगा कि कार सुरक्षित है यदि वह वास्तव में खतरनाक है), लेकिन जैसे-जैसे पहेली गहरी होती जाती है, उत्तर थोड़ा "धुंधला" या कम सटीक हो जाता है। यह क्षितिज को देखकर पहाड़ की दूरी का अनुमान लगाने जैसा है, न कि उसे सटीक रूप से मापने जैसा।

2. "साझा लाइब्रेरी" वाला दृष्टिकोण (फुल्ली शार्डेड डेटा पैरेललिज्म - FSDP)

विचार: एक ऐसी लाइब्रेरी की कल्पना करें जहाँ किताबें इतनी बड़ी हैं कि एक शेल्फ पर फिट नहीं हो सकतीं। हर पाठक के लिए पूरी किताब की कॉपी बनाने के बजाय, लाइब्रेरी किताब को पन्नों में विभाजित करती है।

यह कैसे काम करता है: शोधकर्ता "वेट्स" (किताब के पन्ने) को GPUs में विभाजित करते हैं।
जादुई ट्रिक: जब किसी कंप्यूटर को गणना करने की आवश्यकता होती है, तो वह अन्य कंप्यूटरों से उन सभी पन्नों को जल्दी से इकट्ठा करता है जिनकी उसे आवश्यकता है, गणित करता है, और फिर तुरंत उन पन्नों को वापस रख देता है। किसी भी एक क्षण में, कोई भी कंप्यूटर पूरी किताब अपने पास नहीं रखता।
अच्छी खबर:
- पूर्ण सटीकता: क्योंकि गणित ठीक उसी तरह से किया जाता है जैसे कि एक कंप्यूटर के पास पूरी किताब हो, इसलिए परिणाम सिंगल-कंप्यूटर संस्करण के समान बिट-दर-बिट (bit-for-bit) सटीक होता है। कोई "धुंधलापन" नहीं।
- मेमोरी की बचत: यह बहुत अधिक मेमोरी बचाता है (बेस सेटअप के लिए 80-90%, और पीक उपयोग के लिए 34-39%)।
चुनौती: पन्नों को इकट्ठा करने के लिए कंप्यूटरों के बीच थोड़ी "बातचीत" की आवश्यकता होती है, जिसमें थोड़ा समय लगता है, लेकिन मेमोरी की बचत इसके लायक है।

बड़ी हैरानी: वास्तव में मेमोरी को क्या जाम कर रहा है?

शोधकर्ताओं को उम्मीद थी कि "वेट्स" (पहेली के टुकड़े या किताब के पन्ने) मुख्य समस्या होंगे। वे गलत थे।

एक बार जब उन्होंने वेट्स के लिए जगह खाली करने के लिए इन नए तरीकों का उपयोग किया, तो उन्होंने पाया कि असली बाधा (bottleneck) एक विशेष प्रकार का डेटा है जिसे "अल्फा टेंसर" (alpha tensors) कहा जाता है।

उपमा: कल्पना कीजिए कि आप पहेली सुलझा रहे हैं। "वेट्स" पहेली के टुकड़े हैं, लेकिन "अल्फा टेंसर" वे स्टिकी नोट्स हैं जिन्हें आपको अपनी प्रगति को ट्रैक करने के लिए हर एक टुकड़े पर लिखना पड़ता है।
निष्कर्ष: सबसे उन्नत वेरिफिकेशन मोड में (जहाँ वे 'ब्रांच-एंड-बाउंड' नामक विधि का उपयोग करके दुर्घटनाओं की जाँच करते हैं), ये स्टिकी नोट्स पहेली के टुकड़ों के बजाय 99% मेमोरी लेते हैं।
निष्कर्ष: भले ही उन्होंने पहेली के टुकड़ों को कंप्यूटरों के बीच सफलतापूर्वक विभाजित कर दिया, लेकिन "स्टिकी नोट्स" अभी भी बहुत बड़े हैं। सबसे बड़ी समस्याओं (जैसे सेल्फ-ड्राइविंग कारों के लिए जटिल AI को सत्यापित करना) को हल करने के लिए, भविष्य के काम को यह पता लगाना होगा कि उन स्टिकी नोट्स को भी कंप्यूटरों के बीच कैसे विभाजित किया जाए।

परिणामों का सारांश

टेंसर पैरेललिज्म (Tensor Parallelism): मेमोरी बचाने के लिए बेहतरीन है, लेकिन यह गहरे न्यूरल नेटवर्क के लिए उत्तर को थोड़ा कम सटीक बना देता है।
FSDP: उत्तर को पूरी तरह से सटीक रखता है और बहुत सारी मेमोरी बचाता है। इसने सफलतापूर्वक एक जटिल इमेज-रिकग्निशन मॉडल (ResNet) को सत्यापित किया जिसे पहले चेक करना बहुत बड़ा था।
भविष्य: यहाँ तक कि बड़े AI सिस्टम को सत्यापित करने की कुंजी अब केवल वेट्स को विभाजित करना नहीं है; बल्कि उन "स्टिकी नोट्स" (अल्फा टेंसर) को विभाजित करने का तरीका खोजना है जो सत्यापन प्रक्रिया को ट्रैक करते हैं।

संक्षेप में, यह शोध पत्र दिखाता है कि AI सुरक्षा को सत्यापित करने के लिए कई कंप्यूटरों का उपयोग कैसे किया जाए, लेकिन यह यह भी बताता है कि सबसे बड़े, सबसे जटिल AI सिस्टम को सत्यापित करने से पहले हमें अभी भी एक बड़ी मेमोरी बाधा को पार करना बाकी है।

तकनीकी सारांश: टेंसर पैरेललिज्म और फुली शार्डेड डेटा पैरेललिज्म के साथ न्यूरल नेटवर्क सत्यापन को स्केल करना

1. समस्या विवरण

औपचारिक न्यूरल नेटवर्क सत्यापन (Formal neural network verification) का लक्ष्य गणितीय रूप से यह सिद्ध करना है कि एक नेटवर्क एक निर्दिष्ट डोमेन के भीतर सभी इनपुटों के लिए सुरक्षा गुणों (safety properties) को संतुष्ट करता है। जबकि IBP, CROWN, और $\alpha$ -CROWN जैसे बाउंड-प्रोपैगेशन एल्गोरिदम इस कार्य के लिए अत्याधुनिक (state-of-the-art) हैं, उनका व्यावहारिक अनुप्रयोग GPU मेमोरी द्वारा गंभीर रूप से बाधित है। मानक कार्यान्वयन के लिए आवश्यक है कि वेट मैट्रिसेस (weight matrices) और लीनियर-रिलैक्सेशन $A$ -मैट्रिसेस पूरी तरह से एक ही एक्सेलेरेटर पर मौजूद हों। बड़े नेटवर्क (जैसे, $h=262,144$ चौड़ाई वाले एक सिंगल हिडन लेयर) के लिए, केवल वेट मैट्रिसेस अकेले float32 में 4 GB से अधिक हो सकते हैं, जिससे सत्यापन के लिए आवश्यक $A$ -मैट्रिसेस, बैचिंग, या इंटरमीडिएट टेंसरों के लिए अपर्याप्त मेमोरी बचती है। यह शोध पत्र बड़े पैमाने के मॉडल प्रशिक्षण से पैरामीटर-शार्डिंग रणनीतियों को अनुकूलित करके इस सिंगल-GPU मेमोरी बाधा को दूर करने की आवश्यकता को संबोधित करता है।

2. कार्यप्रणाली (Methodology)

लेखक दो पैरेललिज्म तकनीकों—टेंसर पैरेललिज्म (TP) और फुली शार्डेड डेटा पैरेललिज्म (FSDP)—को auto LiRima / $\alpha, \beta$ -CROWN सत्यापन फ्रेमवर्क के अनुकूल बनाते हैं।

2.1 टेंसर पैरेललिज्म (TP)

TP एक एकल लेयर के भीतर मैट्रिक्स गुणन (matrix multiplications) को $P$ GPUs में विभाजित करता है।

कार्यान्वयन: लेखकों ने कस्टम BoundLinearTP Col और BoundLinearTP Row ऑपरेटर्स लागू किए हैं। वेट्स को आउटपुट (कॉलम) या इनपुट (रो) आयामों के साथ शार्ड किया जाता है।
तंत्र: एक कॉलम-रो पेयर में, गणना वितरित होती है, जिसके लिए परिणाम को पुनः प्राप्त करने के लिए एक AllReduce ऑपरेशन की आवश्यकता होती है।
सीमा: जब शार्ड किए गए ज़ोन के भीतर (कॉलम और रो लेयर्स के बीच) ReLU एक्टिवेशन होते हैं, तो CROWN के बैकवर्ड पास को रिलैक्सेशन ढलान ( $\lambda_i$ ) निर्धारित करने के लिए इंटरमीडिएट बाउंड्स की आवश्यकता होती है। इन्हें सटीक रूप से कंप्यूट करने के लिए हर चरण में इंटर-GPU संचार की आवश्यकता होगी। लोकैलिटी (locality) बनाए रखने के लिए, लेखक इन इंटरमीडिएट बाउंड्स के लिए सरल इंटरवल बाउंड प्रोपैगेशन (IBP) का उपयोग करते हैं। यह एक "रैपिंग इफेक्ट" (wrapping effect) पेश करता है, जिससे शार्ड किए गए ज़ोन की गहराई बढ़ने के साथ बाउंड टाइटनेस (bound tightness) कम हो जाती है।

2.2 फुली शार्डेड डेटा पैरेललिज्म (FSDP)

FSDP पैरामीटर टेंसरों को $P$ GPUs में समान रूप से शार्ड करता है लेकिन गणना के लिए उन्हें पुनर्गठित (reconstruct) करता है।

कार्यान्वयन: लगभग 100 लाइनों के पायथन के साथ कार्यान्वित, FSDP BoundParams टेंसरों को शेड्स (shards) से बदल देता है। एक लेयर की गणना से पहले, एक AllGather पूर्ण वेट मैट्रिक्स को असेंबल करता है; तुरंत बाद, पूर्ण टेंसर को मुक्त कर दिया जाता है, केवल शार्ड को रखा जाता है।
तंत्र: यह दृष्टिकोण सुनिश्चित करता है कि किसी भी क्षण अधिकतम एक पूर्ण वेट मैट्रिक्स GPU मेमोरी में रहे। महत्वपूर्ण रूप से, FSDP अंकगणितीय ऑपरेशनों के क्रम या प्रकार को नहीं बदलता है; यह केवल वेट्स के स्टोरेज स्थान को बदलता है।
सत्यता (Correctness): क्योंकि अंकगणित सिंगल-GPU बेसलाइन के समान ही रहता है, FSDP बिटवाइज़-आइडेंटिकल बाउंड्स ( $\Delta = 0$ ) उत्पन्न करता है।
एकीकरण: FSDP पूर्ण सत्यापन ( $\beta$ -CROWN + Branch-and-Bound) और कनवल्शनल लेयर्स (BoundConv) के साथ सहजता से एकीकृत होता है।

2.3 प्रयोगात्मक सेटअप

प्रयोग NVIDIA A40 GPUs पर PyTorch 2.x और auto LiRPA लाइब्रेरी का उपयोग करके किए गए। लेखकों ने सिंगल-GPU और डिस्ट्रिब्यूटेड मोड के बीच निष्पक्ष तुलना सुनिश्चित करने के लिए एक bab.force_synchronous फ्लैग पेश किया, जो ऑटोमैटिक बैच रिसाइजिंग और प्रूनिंग को अक्षम करके सिंगल-GPU पाथ में समान वर्कलोड सुनिश्चित करता है।

3. मुख्य योगदान

सत्यापन के लिए TP कार्यान्वयन: लेखकों ने कस्टम ऑपरेटर्स का उपयोग करके auto LiRPA के लिए TP लागू किया, जिससे $P=2$ पर लगभग $2\times$ पीक-मेमोरी कमी प्राप्त हुई। उन्होंने ट्रेड-ऑफ का औपचारिक विश्लेषण किया, यह पुष्टि करते हुए कि जबकि साउंडनेस (soundness) संरक्षित रहती है, IBP प्रतिस्थापन के कारण बाउंड टाइटनेस शार्ड किए गए ज़ोन की संख्या के साथ घटती है।
FSDP कार्यान्वयन: उन्होंने बाउंड प्रोपैगेशन के लिए FSDP को लागू किया, जिससे विभिन्न मॉडल आकारों में 80-90% बेसलाइन मेमोरी कमी और वाइड MLPs पर 34-39% पीक मेमोरी कमी प्राप्त हुई। महत्वपूर्ण रूप से, यह विधि सिंगल-GPU बेसलाइन के बिटवाइज़-आइडेंटिकल बाउंड्स उत्पन्न करती है।
पूर्ण सत्यापन के साथ एकीकरण: FSDP को $\beta$ -CROWN और ब्रांच-एंड-बाउंड (BaB) के साथ सफलतापूर्वक एकीकृत किया गया, जिससे FSDP के तहत CIFAR-100 ResNet-large मॉडल का पहला पूर्ण सत्यापन (UNSAT परिणाम) संभव हुआ।
बॉटलनेक की पहचान: व्यापक प्रयोगों के माध्यम से, लेखकों ने पहचाना कि $\alpha$ -CROWN + BaB मोड में, प्राथमिक मेमोरी बॉटलन वेट मैट्रिसेस नहीं, बल्कि प्रति-न्यूरॉन अल्फा टेंसर (ReLU ढलान) हैं।

4. परिणाम

4.1 मेमोरी में कमी

TP: एक वाइड MLP ( $h=262,144$ ) के लिए लगभग $\approx 1.985$ (सैद्धांतिक $P=2$ के करीब) का पीक मेमोरी रिडक्शन प्राप्त किया।
FSDP:
- बेसलाइन मेमोरी: विभिन्न मॉडल आकारों में 80-90% तक कम हो गई, क्योंकि बाउंड गणना से पहले केवल शेड्स आवंटित किए जाते हैं।
- पीक मेमोरी: वाइड MLPs ( $h \ge 4096$ ) पर 34-39% की कमी। बचत बेसलाइन बचत से कम है क्योंकि अनशार्ड $A$ -मैट्रिसेस (जो चौड़ाई के वर्ग के साथ स्केल होते हैं) मेमोरी में बने रहते हैं।

4.2 बाउंड सटीकता और साउंडनेस

TP:
- एक शार्ड ज़ोन वाले मॉडल (कोई इंटरमीडिएट ReLU नहीं) ने सिंगल-GPU रेफरेंस के साथ मशीन-प्रिसिजन समझौता ( $\sim 10^{-8}$ ) प्राप्त किया।
- कई शार्ड ज़ोन वाले मॉडलों ने IBP प्रतिस्थापन के कारण महत्वपूर्ण विचलन दिखाया (उदाहरण के लिए, 3-ज़ोन मॉडल के लिए $\sim 750$ यूनिट), हालांकि साउंडनेस (एकतरफा ओवर-एप्रोक्सिमेशन) संरक्षित रही।
FSDF:
- सभी परीक्षण किए गए मॉडलों (MLPs और ONNX बेंचमार्क) के लिए IBP और CROWN दोनों के लिए बिल्कुल शून्य विचलन ( $\Delta lb = \Delta ub = 0.00$ ) प्राप्त किया।
- एक ट्रांसफार्मर मॉडल (ViT) और कनवल्शनल ResNets को बिटवाइज़-आइडेंटिकल परिणामों के साथ सफलतापूर्वक सत्यापित किया।

4.3 पूर्ण सत्यापन (BaB)

$\beta$ -CROWN + BaB मोड में, मेमोरी खपत अल्फा टेंसर द्वारा संचालित होती है (उदाहरण के लिए, एक विशिष्ट MNIST नेटवर्क में वेट्स के 1.5 MB के मुकाबले अल्फा टेंसरों के लिए 1.5 GB)।

FSDP के माध्यम से वेट्स को शार्ड करने से इस मोड में नगण्य पीक मेमोरी बचत मिली क्योंकि अल्फा टेंसर शार्ड नहीं किए गए थे।
हालांकि, FSDP ने एक ResNet-large मॉडल (VNN-COMP 2024) के सत्यापन को सफलतापूर्वक सक्षम किया जो अन्यथा कठिन या असंभव होता, जिससे BaB ब्रांचिंग के बिना UNSAT प्रॉपर्टी की पुष्टि हुई।

5. महत्व और सीमाएं

यह शोध पत्र दर्शाता है कि प्रशिक्षण से प्राप्त पैरामीटर-शार्डिंग रणनीतियों को मामूली इंजीनियरिंग प्रयास के साथ औपचारिक सत्यापन में स्थानांतरित किया जा सकता है।

महत्व: FSDP सटीक बाउंड्स के साथ बड़े नेटवर्क को सत्यापित करने का मार्ग प्रदान करता है, जबकि TP गहरे नेटवर्क के लिए बाउंड टाइटनेस की कीमत पर उच्चतम मेमोरी रिडक्शन प्रदान करता है। यह कार्य स्थापित करता है कि FSDP जटिल आवश्यकताओं (CNNs, Transformers) के साथ संगत है।
सीमाएं:
- स्केलेबिलिटी: प्रयोग $P=2$ तक सीमित थे। लेखक नोट करते हैं कि जैसे-जैसे $P$ बढ़ता है, अनशार्ड $A$ -मैट्रिसेस संभवतः प्रमुख लागत बन जाएंगे, जिससे घटते रिटर्न (diminishing returns) मिलेंगे।
- TP टाइटनेस: गहरे नेटवर्क के लिए TP में बाउंड टाइटनेस का ह्रास सटीक सत्यापन के लिए इसकी उपयोगिता को सीमित करता है।
- बॉटलनेक शिफ्ट: प्रयोगों से पता चला कि पूर्ण सत्यापन में महत्वपूर्ण स्केलिंग प्राप्त करने के लिए भविष्य के मल्टी-जीपीयू सत्यापन अनुसंधान को केवल वेट्स ही नहीं, बल्कि प्रति-न्यूरॉन अल्फा टेंसरों को शार्ड करने पर ध्यान केंद्रित करना चाहिए।

लेखक निष्कर्ष निकालते हैं कि जबकि FSDF वेट-स्टोरेज बॉटलन को हल करता है, अल्फा टेंसर बॉटलन पूर्ण सत्यापन को बड़े मॉडलों तक स्केल करने के लिए एक महत्वपूर्ण चुनौती बनी हुई है।