PyTorch में OOM को हराना: साधारण GPU पर विशाल ग्राफ़ प्रशिक्षण
Disk Sparse Adam
एक डेवलपर ने Disk Sparse Adam (DSA) बनाया है, जो PyTorch के लिए एक आउट-ऑफ-कोर ऑप्टिमाइज़र है जो mmap के माध्यम से ऑप्टिमाइज़र मोमेंट स्टेट्स को डिस्क पर ले जाता है, जिससे बड़े स्पार्स मॉडल्स के प्रशिक्षण के लिए मेमोरी उपयोग में भारी कमी आती है। यह उपभोक्ता GPU और यहां तक कि मुफ्त Google Colab पर प्रशिक्षण की अनुमति देता है, और बेंचमार्क एक मिलियन इकाइयों पर शून्य VRAM ओवरहेड दिखाता है।
यह लेख डिस्क स्पार्स एडम (DSA) प्रस्तुत करता है, जो PyTorch के लिए एक नया आउट-ऑफ-कोर ऑप्टिमाइज़र है, जिसे उपभोक्ता-ग्रेड GPUs पर लाखों नोड्स वाले विशाल स्पार्स मॉडलों, जैसे कि नॉलेज ग्राफ एम्बेडिंग, को प्रशिक्षित करने के लिए डिज़ाइन किया गया है। मानक torch.optim.SparseAdam पहले और दूसरे क्षण की स्थितियों को मेमोरी में संग्रहीत करता है, जो 10 मिलियन इकाइयों और 128-आयामी वैक्टर वाले मॉडल के लिए लगभग 10.24 GB अतिरिक्त मेमोरी की खपत करता है, जिसके परिणामस्वरूप अक्सर CUDA आउट-ऑफ-मेमोरी त्रुटियाँ होती हैं। DSA इसके बजाय इन क्षण मैट्रिसेस को बाइनरी फ़ाइलों के रूप में डिस्क पर संग्रहीत करता है और उन्हें OS mmap तंत्र का उपयोग करके मेमोरी में मैप करता है, प्रत्येक अनुकूलन चरण पर केवल सक्रिय बैच की स्थितियों को पढ़ता और लिखता है। ऑप्टिमाइज़र को मानक PyTorch पाइपलाइनों के लिए ड्रॉप-इन प्रतिस्थापन के रूप में डिज़ाइन किया गया है, और यह यूक्लिडियन और हाइपरबोलिक (पॉइन्केयर बॉल) दोनों स्थानों का समर्थन करता है। लेखक रिपोर्ट करते हैं कि ऑप्टिमाइज़र राज्यों के लिए मेमोरी उपयोग गीगाबाइट्स से घटकर मेगाबाइट्स हो जाता है, और 1 मिलियन इकाइयों के साथ Kaggle पर एक बेंचमार्क में 0.00 MB VRAM ओवरहेड और 134,212 नमूने प्रति सेकंड की थ्रूपुट दिखी। DSA GitHub पर MIT लाइसेंस के तहत उपलब्ध है, जिसमें सीमाएँ शामिल हैं जैसे कि सर्वोत्तम प्रदर्शन के लिए NVMe SSD की आवश्यकता और केवल स्पार्स ग्रेडिएंट्स के साथ संगतता, जैसे कि torch.nn.Embedding या EmbeddingBag से प्राप्त।
स्रोत: Habr — хаб ИИ —
मूल