TL;DR ML मॉडल लोडिंग धीमी है, भले ही Linux पेज कैश गर्म हो। इसलिए हमने इसे तेज बनाने के लिए एक लाइब्रेरी बनाई। कुछ दिलचस्प तकनीकी विवरण हैं जिन्हें हम साझा करना चाहते हैं, इसलिए हमने यह ब्लॉग लिखा। लाइब्रेरी ने एक अप्रत्याशित प्रभाव भी डाला, जिसे अंत में चर्चा की गई है।
तर्क
यह सब 2 साल पहले शुरू होता है, जब हमने लो-पॉली जनरेशन मोड का पहला प्रयास जारी किया। लो-पॉली मोड अच्छा नहीं गया, यह आज के दृष्टिकोण से खराब परिणाम देता है, लेकिन हमने इसके लिए बहुत भुगतान किया - एक समर्पित GPU केवल एक दिन में एकल अंक कार्यों को संसाधित करता है। इसमें फाइन-ट्यून किए गए वेट्स हैं, जो सभी अन्य मॉडल वेट्स को VRAM से बाहर धकेलने के लिए पर्याप्त बड़े हैं। इससे भी बदतर, हमारे पास शायद 3 ऐसे मॉडल हैं (सटीक संख्या याद नहीं है), उन्होंने हमारे इन्फरेंस इंफ्रा का एक महत्वपूर्ण हिस्सा बनाया, जिससे एक काफी कठोर दक्षता अनुपात बन गया। और नहीं, हम मॉडल्स को बस-इन-टाइम लोड नहीं कर सकते, इसमें 30 सेकंड लगते हैं, जो वास्तविक प्रोसेसिंग समय से अधिक है।
तब हमारे पास समर्पित पाइपलाइन इंजीनियर नहीं थे, हमारे एल्गोरिदम डेवलपर्स ने इसे वर्कअराउंड करने की पूरी कोशिश की। कुछ दिनों बाद, हमारा कोडबेस this.to('cpu') और that.to('cuda') से भरा हुआ था। यह दृष्टिकोण कुछ समय के लिए काम करता है, लेकिन समय-समय पर हमारे एल्गो डेवलपर्स के प्रवाह को तोड़ देता है। क्या होगा अगर चीजें स्वचालित रूप से हो सकती हैं? यह पायथन है, चीजें पायथन में स्वचालित रूप से होती हैं।
आप 'स्वचालित रूप से' कैसे परिभाषित करते हैं?
आइए एक एल्गोरिदम डेवलपर की भूमिका में कूदें। चीजें काफी स्पष्ट हैं: मैं अपने मुख्य एल्गोरिदम के बाहर के रनटाइम प्रदर्शन की परवाह नहीं करना चाहता जब तक कि मुझे बिल्कुल करना न पड़े। मैं मॉडल स्वैप इन और आउट के बारे में कुछ भी नहीं जानना चाहूंगा।
बेशक हम इसे हासिल नहीं कर सकते, लेकिन हम एल्गोरिदम कोड में हस्तक्षेप को कम करने की कोशिश कर सकते हैं। यह मुझे gevent लाइब्रेरी की मंकी-पैचिंग की याद दिलाता है, यह (मुख्य रूप से) socket लाइब्रेरी को पैच करता है, इसे gevent.socket से बदल देता है जो IO ब्लॉक होने पर अन्य ग्रीनलेट्स पर स्विच कर सकता है, बहुत कुछ एक गोरूटीन की तरह (वास्तव में gevent Golang से पुराना है!)।
चूंकि हम उस समय केवल HuggingFace लाइब्रेरीज़ (transformers, diffusers) का उपयोग मॉडल लोड करने के लिए कर रहे थे, लक्ष्य स्पष्ट हो गया: हम केवल एक मंकी-पैच कॉल पेश करते हैं, और बाकी कोड अपरिवर्तित रहना चाहिए, XXXPipeline.from_pretrained(...) बहुत तेज होना चाहिए।
कुछ तथ्य, स्पष्ट निर्णय और धारणाएँ
Overmind एक कैशिंग लाइब्रेरी है, यह मॉडल लोडिंग कॉल के परिणामों को सिस्टम मेमोरी में कैश करता है और बाद में इसे तेजी से पुनर्निर्मित करता है।
हम मंकी-पैचिंग कैसे लागू की जाती है, इस पर चर्चा करना छोड़ देते हैं, यह एक इतना दिलचस्प विवरण नहीं है। हमें केवल यह जानने की जरूरत है कि यह सभी XXXPipeline.from_pretrained(...) कॉल्स को overmind.api.load(XXXPipeline.from_pretrained, ...) पर पुनर्निर्देशित करता है।
हम अपने कैश परिणाम को सीरियलाइज़ करने के लिए pickle का उपयोग करते हैं क्योंकि... हमारे पास कोई विकल्प नहीं है, और torch.save स्वयं pickle का उपयोग करता है, इसका उपयोग न करना अजीब है।
हम एक क्लाइंट/सर्वर आर्किटेक्चर का उपयोग करते हैं क्योंकि हम नहीं चाहते कि हमारा कैश प्रक्रिया समाप्त होने पर अमान्य हो जाए। कई सबप्रोसेस कॉल्स इससे लाभान्वित हो सकते हैं।
हम मानते हैं कि XXXPipeline.from_pretrained पैरामीटर सरल हैशेबल चीजें (str और इसी तरह की चीजें) हैं और अन्य मॉडल overmind द्वारा लोड किए गए हैं (बाद में समझाया गया)।
नाम overmind Starcraft से उधार लिया गया है, जैसा कि आपने अनुमान लगाया होगा।
इसे तेजी से पुनर्निर्मित करें!
हम pickle.loads परिणाम को मेमोरी में सहेजकर और इसे एक दिन नहीं कह सकते। आखिरकार, एक गर्म किए गए परिदृश्य में, Linux पेज कैश ने ऑन-डिस्क मॉडल्स को कैश करने का काम किया और हम अभी भी लोडिंग समय को दसियों सेकंड में माप सकते हैं।
अप्रभावशीलता मेमोरी कॉपी से आती है। पायथन में, लाखों ऑब्जेक्ट्स बनाना भी कुछ सौ मिलीसेकंड से अधिक नहीं लगेगा। हालांकि, 10GiB की मेमोरी कॉपी के लिए, इसमें आधा सेकंड लगेगा। हमें मेमोरी कॉपी से यथासंभव बचना चाहिए।
सौभाग्य से, अधिकांश बड़े मेमोरी चंक्स Torch टेन्सर्स हैं, हम सुरक्षित रूप से केवल उन्हें संबोधित कर सकते हैं और बाकी को अनदेखा कर सकते हैं।
वास्तव में, मुझे टेन्सर शेयरिंग तंत्र पर शोध करते समय Torch टेन्सर की आंतरिक संरचना का ज्ञान प्राप्त हुआ:
# torch.multiprocessing.reductions से कॉपी किया गया, अधिकांश कोड हटा दिया गया है
def reduce_tensor(tensor):
...
storage = tensor._typed_storage()
...
metadata = (
tensor.storage_offset(),
tensor.size(),
tensor.stride(),
tensor.requires_grad,
)
return (rebuild_tensor, (type(tensor), storage, metadata))काफी सरल: एक टेन्सर उसका प्रकार, उसका मेटाडेटा और उसका अंतर्निहित स्टोरेज होता है। यहाँ storage का प्रकार TypedStorage है, लेकिन वास्तव में TypedStorage सिर्फ UntypedStorage का एक सरल रैपर है। UntypedStorage वह क्लास है जो वास्तव में सभी टेन्सर डेटा को धारण करता है।
हमारा कार्य अब अधिक विशिष्ट हो गया है: हम UntypedStorage की कॉपी करने से कैसे बच सकते हैं? क्या हम इन टेन्सर मेमोरी को स्वयं प्रबंधित कर सकते हैं और UntypedStorages को उस मेमोरी की ओर इशारा करके बना सकते हैं जिसे हम प्रबंधित करते हैं?
उत्तर है हाँ!
जहाँ UntypedStorage का निर्माण होता है, उसके C++ कोड को सरसरी निगाह से देखने पर, हम आसानी से ऐसा कोड स्निपेट पा सकते हैं:
// torch/csrc/Storage.cpp से कॉपी किया गया
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
// ...असंबंधित कोड को छोड़ते हुए...
auto new_storage_impl = make_storage_impl(
c10::StorageImpl::use_byte_size_t(),
slicelength,
at::DataPtr(
static_cast<void*>(data + start),
old_storage_impl,
[](void* s) {
c10::raw::intrusive_ptr::decref(static_cast<at::StorageImpl*>(s));
},
old_storage_impl->device()),
old_storage_impl->allocator(),
/* resizable */ false,
device_opt);
PyObject* _ret =
THPStorage_NewWithStorage(Py_TYPE(self), std::move(new_storage_impl));
return _ret;
}हम न केवल एक पॉइंटर का उपयोग कर सकते हैं, बल्कि at::DataPtr क्लास विनाश को भी संभाल सकता है, जिससे जीवनकाल प्रबंधन बहुत सरल हो जाता है।
Python पक्ष पर, एक मेमोरी क्षेत्र के लिए पॉइंटर को memoryview ऑब्जेक्ट द्वारा दर्शाया जाता है, ये ऑब्जेक्ट बफर प्रोटोकॉल का समर्थन करते हैं। हम कई चीजों से memoryview ऑब्जेक्ट प्राप्त कर सकते हैं, bytes और mmap वे दो मुख्य चीजें हैं जो इसे समर्थन करती हैं, और वे भी हैं जिनकी हमें परवाह है।
अंत में, हम जानते हैं कि हमें क्या करना चाहिए: एक फ़ंक्शन बनाएं जो एक memoryview ऑब्जेक्ट को स्वीकार करता है और उसे बिना कॉपी किए UntypedStorage में बदल देता है। memoryview से UntypedStorage को पुनर्निर्मित करने की क्षमता के साथ, वास्तविक टेन्सर डेटा को पिकल स्ट्रीम में नहीं होना पड़ता है, जिससे हमें कॉपी करने के लिए डेटा का आकार काफी कम हो जाता है।
void initOvermindHelpers(py::module m) {
// ...
m.def("_make_untyped_storage", [](py::buffer b) {
auto info = new py::buffer_info(b.request());
return pybind11::reinterpret_steal<py::object>(THPStorage_NewWithStorage(
THPStorageClass,
c10::make_intrusive<at::StorageImpl>(
c10::StorageImpl::use_byte_size_t(),
info->size,
at::DataPtr(
info->ptr, info,
[](void* ptr) {
py::gil_scoped_acquire gil;
auto b = static_cast<py::buffer_info*>(ptr);
delete b;
},
at::DeviceType::CPU
),
/*allocator=*/nullptr,
/*resizable=*/false,
)
));
});
}यह overmind का मुख्य निर्माण खंड है।
टेन्सर साझा करना!
नोट: PyTorch में पहले से ही एक टेन्सर साझा करने की प्रणाली है, लेकिन यह हमारी आवश्यकताओं को पूरा नहीं करती। इस पर बाद में और अधिक।
पहले, क्लाइंट और सर्वर के बीच मेमोरी साझा करना
जब हम 'साझा' और 'मेमोरी' को एक साथ देखते हैं, तो हम सभी को shmget और उसके दोस्तों का उपयोग करने की इच्छा होती है। यह एक मेमोरी साझा करने की प्रणाली के रूप में "डिज़ाइन" किया गया है, है ना? लेकिन इसके 2 प्रमुख दोष हैं:
- POSIX shm एक दुर्लभ संसाधन है, आप जो उपयोग कर सकते हैं वह इस बात पर निर्भर करता है कि सिस्टम एडमिन ने सिस्टम को कैसे कॉन्फ़िगर किया है। एक चरम लेकिन सर्वव्यापी उदाहरण है Docker कंटेनर, डिफ़ॉल्ट रूप से आपके पास केवल 64MiB POSIX shm उपयोग योग्य है।
- POSIX shm आपके प्रोसेस से अधिक समय तक जीवित रहता है, आपको अपनी प्रबंधन प्रक्रिया स्वयं करनी होती है। यदि प्रबंधन प्रक्रिया को जबरन समाप्त कर दिया जाता है, या इसे सावधानीपूर्वक नहीं संभाला गया, तो shm ऑब्जेक्ट सिस्टम पर अनिश्चित काल तक रह सकता है।
यदि आप ध्यान से देखें, तो Linux में कई दिलचस्प सिस्टम कॉल्स होते हैं। memfd_create वह है जिसमें हम रुचि रखते हैं: यह आपको एक fd देता है जो गुमनाम मेमोरी के आवंटन का प्रतिनिधित्व करता है। आप इस पर सभी प्रकार के फ़ाइल संचालन कर सकते हैं: पढ़ना, लिखना, और, निश्चित रूप से, mmap। यदि हम fd साझा कर सकते हैं, तो हम मेमोरी साझा कर सकते हैं।
एक fd साझा करने का एक 'मानक' लेकिन रहस्यमय तरीका है: sendmsg के साथ SCM_RIGHTS। हम लाइब्रेरीज़ का लाभ उठा सकते हैं ताकि sendmsg प्रक्रिया के डरावने विवरणों को छिपा सकें, लेकिन हमें अभी भी सर्वर और क्लाइंट प्रक्रियाओं के बीच समन्वय करना होगा। हमने यहां एक हैक का उपयोग करने का निर्णय लिया: क्लाइंट साइड पर /proc/{pidof(server)}/fd/{memfd} को बस खोलें, जबकि overmind सर्वर साइड पर fd को कभी बंद न करें। केवल आवश्यक संचार एक (pid, fd) ट्यूपल है। यह हमारे मामले में पूरी तरह से काम करता है।
उपरोक्त शब्द इन लाइनों में संक्षेपित होते हैं:
class SharedMemory:
@classmethod
def create(cls, shift):
# सर्वर साइड पर कॉल किया गया
libc = ctypes.CDLL(None)
name = _make_filename(shift).encode('utf-8')
fd = libc.memfd_create(name, os.O_RDWR)
os.ftruncate(fd, 1 << shift)
mem_id = (os.getpid(), fd)
return cls(fd=fd, mem_id=mem_id)
@classmethod
def rebuild(cls, mem_id):
# क्लाइंट साइड पर कॉल किया गया
pid, fd = mem_id
local_fd = os.open(f'/proc/{pid}/fd/{fd}', os.O_RDWR)
return cls(fd=local_fd, mem_id=mem_id)
def get_buffer(self):
# दोनों साइड पर कॉल किया गया
self._mmap = mmap.mmap(self._fd, size)
self._buf = memoryview(self._mmap)
return self._bufपिकलिंग के साथ एकीकृत करें
जैसा कि हमने पहले चर्चा की, हमें UntypedStorage की पिकलिंग प्रक्रिया को संशोधित करने की आवश्यकता है। जैसा कि torch.multiprocessing.reductions में लागू किया गया था, हम pickle के लिए अपने कस्टम रिड्यूस फ़ंक्शंस को परिभाषित करते हैं:
# Hoarder और borrower ऊपर SharedMemory के लिए एक रैपर है, जिसमें
# बोरिंग चीजें जैसे मेमोरी एरीना, आदि शामिल हैं।
def _reduce_storage(storage):
# सर्वर द्वारा कॉल किया गया
device = storage.device
storage = storage.cpu()
# साझा मेमोरी में सामग्री संग्रहीत करें
# `frag` में सामग्री को खोजने के लिए आवश्यक पूरी जानकारी होती है।
frag = hoarder.put(storage)
return (_rebuild_storage_on_client, (frag, device))
def _rebuild_storage_on_client(frag, device):
# क्लाइंट द्वारा कॉल किया गया
mv = borrower.borrow(frag) # साझा मेमोरी से एक memoryview प्राप्त करें
storage = _make_untyped_storage(mv) # जीरो-कॉपी!
if device.type == 'cuda':
return storage.cuda(device.index)
return storage
class OvermindPickler(dill.Pickler):
...
OvermindPickler.register(torch.storage.UntypedStorage, _reduce_storage)अब, सरल OvermindPickler.dumps और OvermindPickler.loads साझा मेमोरी का उपयोग करके गति बढ़ाएंगे। यदि आप पहले से ही थक चुके हैं, तो आप यहां पढ़ना बंद कर सकते हैं। बाकी विवरण हैं।
विवरण में शैतान
क्यों नहीं PyTorch का इन-हाउस टेंसर शेयरिंग?
'इन-हाउस टेंसर शेयरिंग' से मेरा मतलब है torch.multiprocessing.reductions।
- उच्च स्तर पर, PyTorch की विधियाँ 'टेंसर को सबप्रोसेस में पास करने' के लिए डिज़ाइन की गई हैं, जो समान लगती हैं लेकिन सूक्ष्म अंतर के साथ।
- PyTorch मेमोरी साझा करने के लिए POSIX shm का उपयोग करता है, जो पहले उल्लेखित सीमा के अधीन है।
- प्रत्येक टेंसर (या
UntypedStorage) के लिए, PyTorch इसके लिए एक समर्पित POSIX shm ऑब्जेक्ट आवंटित करता है, भले ही उसमें केवल 4 बाइट्स हों। प्रत्येक ऑब्जेक्ट एक fd का उपभोग करता है। - PyTorch POSIX shm को डीएलोकेट करता है जब वे अनपिकल्ड होते हैं, जो हमारी आवश्यकताओं के लिए अनुपयुक्त बनाता है। हमें एक ही पिकल स्ट्रीम को कई बार डीसिरियलाइज़ करने की आवश्यकता है।
- बहुत सारी CUDA संबंधित शेयरिंग लॉजिक हैं, जो हमारे उपयोग के मामले के लिए शुद्ध शोर और परेशानी हैं।
आप क्यों कहते हैं 'टेंसर डेटा कई बार कॉपी किया गया'?
एक सामान्य ऑन-डिस्क torch.load के लिए:
- ऑन-डिस्क
torch.saveफ़ाइल को मेमोरी में पढ़ा जाता है। - वास्तविक
torch.UntypedStorageडेटा कोbytesके रूप में Zip फ़ाइल निष्कर्षण द्वारा प्राप्त करें (torch.saveएक ज़िप फ़ाइल उत्पन्न करता है)। - C++ कोड अपने प्रबंधित मेमोरी में
torch.UntypedStorageकंस्ट्रक्टर में डेटा की प्रतिलिपि बनाएगा।
एक साधारण pickle.dumps और बाद में pickle.loads के लिए:
- उत्पन्न पिकल स्ट्रीम आंतरिक रूप से एक और पिकल स्ट्रीम को एम्बेड करता है,
pickle.loadsआंतरिक स्ट्रीम को एक नएbytesमें कॉपी करेगा। torch.UntypedStorageडेटा आंतरिक पिकल स्ट्रीम में एम्बेड होता है,torch.UntypedStorageके निर्माण के समय एक और प्रतिलिपि होती है।- C++ कोड अपने प्रबंधित मेमोरी में
torch.UntypedStorageकंस्ट्रक्टर में डेटा की प्रतिलिपि बनाएगा।
diffusers के पास एक डायनामिक मॉड्यूल है
मॉडल रिपोज़ में Python फ़ाइलें शामिल हो सकती हैं जो रनटाइम पर diffusers_modules namespace में आयात की जाती हैं। क्लाइंट के पास ये sys.path में नहीं होते, जिससे अनपिकलिंग टूट जाती है। सौभाग्य से, diffusers इन डायनामिक Python फ़ाइलों को डिस्क पर लिख देगा, इसलिए हम बस मॉड्यूल आयात कर सकते हैं और इसे एक दिन कह सकते हैं।
def diffusers_dyn_module_workaround():
from diffusers.utils.constants import HF_MODULES_CACHE
modpath = Path(HF_MODULES_CACHE) / "diffusers_modules/__init__.py"
spec = importlib.util.spec_from_file_location("diffusers_modules", modpath)
sys.modules["diffusers_modules"] = importlib.util.module_from_spec(spec)bitsandbytes के लिए समर्थन
bitsandbytes का समर्थन करने में सबसे कष्टप्रद बात यह है कि क्वांटाइजेशन प्रक्रिया GPU पर होती है। एक बार जब हमने overmind सर्वर में CUDA और torch को इनिशियलाइज़ कर लिया, तो इसे अनइनिशियलाइज़ करने का कोई आसान तरीका नहीं है, जो वास्तविक वर्कलोड्स के लिए समस्याएं पैदा कर सकता है (मुख्य रूप से कम उपयोगी VRAM)। इसलिए, हमने अपने सर्वर को एक सबप्रोसेस स्पॉन करने, इसे साझा मेमोरी में लोड करने और समाप्त करने के लिए संशोधित किया। इससे overmind सर्वर की स्थिरता में सुधार होता है।
क्वांटाइज्ड पैरामीटर bitsandbytes द्वारा प्रदान किए गए विशेष सबक्लास हैं। उन्हें 'पिकलिबिलिटी' के लिए डिज़ाइन नहीं किया गया था, इसलिए हमें इसे स्वयं करना होगा।
def _reduce_bnb_param(p):
dev = p._prev_device
assert p.quant_state
return (_rebuild_bnb_param, (type(p), p.data, p.quant_state.as_dict(packed=True), dev))
def _rebuild_bnb_param(typ, data, qs_dict, dev):
return typ.from_prequantized(data, qs_dict, device=dev)
def bitsandbytes_quirks():
try:
import bitsandbytes
except ImportError as e:
return
ForkingPickler.register(bitsandbytes.nn.modules.Params4bit, _reduce_bnb_param)
ForkingPickler.register(bitsandbytes.nn.modules.Int8Params, _reduce_bnb_param)bitsandbytes के माध्यम से क्वांटाइज्ड मॉडल हुक और मंकी-पैच के साथ आते हैं जो पिकल नहीं होते, हमें उन्हें हटाना होगा:
from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None) # चेतावनी मंकीपैच हटाएं
model.__dict__.pop('cuda', None)हमने ऐसे मुद्दों का भी सामना किया है जहां फ़ंक्शंस अन्य फ़ंक्शंस के भीतर नेस्टेड होते हैं (बल्कि शीर्ष स्तर पर होने के बजाय), जो उन्हें पिकल नहीं बनाता है। हमने इसे वर्कअराउंड करने की कोशिश की, लेकिन कोई सफलता नहीं मिली। हमें अपने पिकल को stdlib द्वारा प्रदान किए गए से dill में स्विच करना पड़ा ताकि इसे पिकल किया जा सके। dill बहुत अधिक शक्तिशाली है, लेकिन यह एक शुद्ध Python कार्यान्वयन है, जो मानक लाइब्रेरी संस्करण की तुलना में बहुत धीमा है। सौभाग्य से, यह लागत केवल एक बार चुकाई जाएगी जब हम मॉडल को पहली बार लोड कर रहे हैं (केवल पिकलिंग को प्रभावित करता है, अनपिकलिंग को नहीं)।
stable-fast के लिए समर्थन
stable-fast torch.compile परिणाम उत्पन्न करता है, जिन्हें पिकल नहीं किया जा सकता। लेकिन torch.jit.save के साथ, हम परिणामों को एक ज़िप फ़ाइल के रूप में सहेज सकते हैं। यह अक्षम लगता है, लेकिन कुछ न होने से बेहतर है।
केवल torch.jit.save के साथ stable-fast परिणामों को पिकल करना पर्याप्त नहीं है। stable-fast Torch मॉड्यूल को ट्रेस करने योग्य बनाने के लिए एक 'फ्लैटन' प्रक्रिया का उपयोग करता है। जब कुछ ऐसा मिलता है जिसे वह पहचान नहीं पाता (उदाहरण के लिए, dataclass की क्लास), तो वह इसे सीरियलाइज़ नहीं करेगा, बल्कि केवल वास्तविक क्लास के संदर्भ को रखेगा। हमने संबंधित लॉजिक को पैच किया है ताकि 'फ्लैटन' की गई स्ट्रीम के भीतर एक पिकल की गई क्लास को वास्तव में स्टोर किया जा सके।
def stable_fast_quirks():
...
# pickle dataclass type instead of just put it into a container (which will not survive after torch.jit.save)
def flatten_dataclass(obj):
from sfast.utils.flat_tensors import flatten_bytes, flatten_dict
import dataclasses
d = dict((field.name, getattr(obj, field.name))
for field in dataclasses.fields(obj))
import pickle
pickled = pickle.dumps(obj.__class__)
return flatten_bytes(pickled) + flatten_dict(d)
def unflatten_dataclass(tensors, start):
from sfast.utils.flat_tensors import unflatten_bytes, unflatten_dict
import pickle
pickled, start = unflatten_bytes(tensors, start)
clz = pickle.loads(pickled)
content, start = unflatten_dict(tensors, start)
return clz(**content), start
sfast.utils.flat_tensors.flatten_dataclass = flatten_dataclass
sfast.utils.flat_tensors.unflatten_dataclass = unflatten_dataclassयहाँ दो और ट्रिक्स हैं:
- हम ZIP फ़ाइल को
ZIP_STOREDके साथ पुनः पैक करते हैं, ताकि हमें हर बार लोड करने पर ZIP फ़ाइल को डीकंप्रेस करने की आवश्यकता न हो। torch.jit.loadइंटरफ़ेस भी मेमोरी कॉपी समस्या का कारण बनता है, इसलिए हमने इसे Python बफर प्रोटोकॉल के माध्यम से लोड करने के लिए एक सरल रैपर लिखा, जैसेUntypedStorage।
void initOvermindHelpers(py::module m) {
// ...
m.def("import_ir_module_from_buffer_0copy",
[](std::shared_ptr<torch::jit::CompilationUnit> cu, py::buffer buffer) {
auto info = buffer.request();
imemstream in((char*)info.ptr, info.size); // No copy!
return import_ir_module(std::move(cu), in, ...);
}
);
}vae=vae पैटर्न
हमारे कोडबेस में कुछ ऐसा है, यह एक मॉडल को लोड करने का प्रयास करता है जो पहले से लोड किए गए मॉडल को अपने तर्क के रूप में लेता है:
import overmind.api
overmind.api.monkey_patch_all()
import torch
from diffusers.models import AutoencoderKL
from diffusers import (
ControlNetModel,
StableDiffusionControlNetPipeline,
)
vae = AutoencoderKL.from_pretrained(
"lemon2431/ChineseInkComicStrip_v10",
subfolder="vae",
torch_dtype=torch.float16,
)
controlnet_depth = ControlNetModel.from_pretrained(
"lllyasviel/control_v11f1p_sd15_depth",
torch_dtype=torch.float16,
variant="fp16",
)
controlnet_edge = ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_softedge",
torch_dtype=torch.float16,
variant="fp16",
)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
"lemon2431/ChineseInkComicStrip_v10",
vae=vae, # यहाँ !
controlnet=[controlnet_edge, controlnet_depth], # और यहाँ!
torch_dtype=torch.float16,
safety_checker=None,
)
pipeline.to('cuda')जैसा कि हमने पहले उल्लेख किया है, फ़ंक्शन तर्कों को सरल, आसानी से पिकल करने योग्य वस्तुओं के रूप में माना जाता है, लेकिन यह पैटर्न उस धारणा को तोड़ता है। इसे संभालने के लिए, हमने विशेष तर्क जोड़ा: प्रत्येक कैश किए गए परिणाम को एक ID संलग्न किया जाता है। यदि वह वस्तु किसी अन्य कॉल में तर्क के रूप में उपयोग की जाती है, तो क्लाइंट इसे उसके ID के साथ बदल देता है, और सर्वर तब ID के आधार पर वास्तविक वस्तु को पुनः प्राप्त कर सकता है।
परिणामी pipeline मॉडल में vae का संदर्भ होगा। सरलता के लिए, हम इसे सीधे पिकल करते हैं। हालाँकि, जब वास्तविक UntypedStorage को साझा मेमोरी में ले जाया जाता है, तो हम किसी भी दोहराए गए डेटा को डेडुप्लिकेट करते हैं।
हमने पिकल के persistent_id मैकेनिज़्म का उपयोग कर सकते थे, लेकिन मैंने इस मार्ग को नहीं आजमाया। यह थोड़ा अफसोसजनक है।
बेंचमार्किंग
और अब वह भाग जिसे हर कोई देखना पसंद करता है।
हम अपने परीक्षण के लिए पिछले खंड के VAE पैटर्न स्क्रिप्ट का उपयोग करते हैं।
| परीक्षण | vae | depth | edge | pipeline | to('cuda') | कुल |
|---|---|---|---|---|---|---|
| बिना, 1st | 1.18 | 0.98 | 1.41 | 1.65 | 0.91 | 6.16 |
| बिना, 2nd | 1.15 | 0.96 | 0.97 | 1.65 | 0.89 | 5.66 |
| बिना, 3rd | 1.15 | 0.96 | 0.98 | 1.61 | 0.91 | 5.65 |
| w/o, 4th | 1.42 | 1.10 | 1.11 | 1.72 | 0.88 | 6.27 |
| w/o, 5th | 1.28 | 1.08 | 1.10 | 1.72 | 0.92 | 6.13 |
| w/, 1st | 5.44 | 5.17 | 5.41 | 7.29 | 0.86 | 24.20 |
| w/, 2nd | 0.00 | 0.01 | 0.01 | 0.20 | 0.87 | 1.12 |
| w/, 3rd | 0.01 | 0.01 | 0.01 | 0.21 | 0.86 | 1.12 |
| w/, 4th | 0.01 | 0.01 | 0.01 | 0.20 | 0.90 | 1.15 |
| w/, 5th | 0.01 | 0.01 | 0.01 | 0.21 | 0.86 | 1.13 |
जैसा कि आप देख सकते हैं, overmind के साथ प्रारंभिक लोडिंग में 24.2 सेकंड लगते हैं, जो इसके बिना लोडिंग की तुलना में काफी लंबा है। हालांकि, बाद के लोड्स पर, केवल .to('cuda') की लागत अभी भी बनी रहती है।
सभी सीरियलाइज्ड मॉडल फाइलों के आकार को जोड़ते हुए, पूरे पाइपलाइन के लिए लगभग 5808 मेगाबाइट मेमोरी का उपयोग होने का अनुमान है। एक त्वरित बेंचमार्क एक समान परिणाम देता है।
In [1]: t = torch.ones((5808, 1024, 1024), dtype=torch.uint8)
In [2]: %time a = t.cuda()
CPU times: user 976 ms, sys: 874 μs, total: 977 ms
Wall time: 976 ms
In [3]: %timeit a = t.cuda()
1.01 s ± 56.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)Intel i9-11900K + GeForce RTX 4090 पर परीक्षण किया गया।
अप्रत्याशित साइड इफेक्ट्स (सकारात्मक!)
हमारा प्राथमिक उद्देश्य overmind का निर्माण करना था ताकि इन्फेरेंस के दौरान मॉडल वेट्स को तेजी से स्विच किया जा सके। जबकि इसने अपने उद्देश्य को पूरा किया, हमने कई अतिरिक्त लाभ भी खोजे।
हम अपने एप्लिकेशन के कई इंस्टेंस तैनात करते हैं, प्रत्येक GPU के लिए एक। इस प्रकार, प्रति नोड 8 प्रक्रियाएँ होंगी। जब हमने overmind को तैनात किया, तो सिस्टम मेमोरी का उपयोग नाटकीय रूप से कम हो गया। हम सिस्टम मेमोरी की कमी से पीड़ित नहीं थे, लेकिन अगर होते, तो यह एक बड़ी जीत होती।
बाद में, हमने इसे हमारे एल्गोरिदम और पाइपलाइन डेवलपर्स के लिए एक बड़ा बढ़ावा पाया। प्रत्येक संशोधन-सत्यापन लूप के लिए, हम लोडिंग समय के 10 से 20 सेकंड बचा सकते थे, जो एक बड़ी संख्या में जुड़ सकता था। अधिक महत्वपूर्ण बात यह है कि बचाए गए सेकंड डेवलपर्स को प्रवाह में बनाए रख सकते थे।
Github
हम इसे Github पर ओपन-सोर्स कर रहे हैं, अगर इससे मदद मिली तो हमें खुशी होगी।
