TL;DR การโหลดโมเดล ML ช้า แม้จะมีการเตรียมแคชหน้าเพจของ Linux ไว้แล้ว ดังนั้นเราจึงสร้างไลบรารีเพื่อทำให้มันเร็วขึ้น มีรายละเอียดทางเทคนิคที่น่าสนใจที่เราอยากแบ่งปัน ดังนั้นเราจึงเขียนบล็อกนี้ ไลบรารียังมีผลกระทบที่ไม่คาดคิด ซึ่งจะกล่าวถึงในตอนท้าย
เหตุผล
ทั้งหมดเริ่มต้นเมื่อ 2 ปีที่แล้ว เมื่อเราส่งโหมดการสร้าง lowpoly เป็นครั้งแรก โหมด lowpoly ไม่เป็นไปด้วยดี มันให้ผลลัพธ์ที่ไม่ดีจากมุมมองของวันนี้ แต่เราจ่ายเงินไปมากสำหรับมัน -- GPU เฉพาะทางประมวลผลงานได้เพียงตัวเลขหลักเดียวต่อวัน มันมีน้ำหนักที่ปรับแต่งมาอย่างดี ใหญ่พอที่จะขับไล่น้ำหนักโมเดลอื่น ๆ ออกจาก VRAM ที่แย่กว่านั้น เรามีโมเดลแบบนี้ประมาณ 3 โมเดล (จำจำนวนที่แน่นอนไม่ได้) พวกมันเป็นส่วนสำคัญของโครงสร้างพื้นฐานการอนุมานของเรา ทำให้อัตราประสิทธิภาพค่อนข้างไม่ให้อภัย และไม่ เราไม่สามารถโหลดโมเดลแบบ just-in-time ได้อย่างง่ายดาย มันใช้เวลา 30 วินาที ซึ่งใหญ่กว่าความจริงที่ใช้ในการประมวลผล
ตอนนั้นเราไม่มีวิศวกรท่อส่งข้อมูลเฉพาะทาง นักพัฒนาการอัลกอริทึมของเราพยายามอย่างดีที่สุดเพื่อหาวิธีแก้ปัญหานี้ หลายวันต่อมา โค้ดเบสของเราถูกทิ้งด้วย this.to('cpu') และ that.to('cuda') วิธีนี้ใช้ได้ผลชั่วคราว แต่ทำให้การทำงานของนักพัฒนาอัลกอริทึมของเราหยุดชะงักเป็นครั้งคราว จะเกิดอะไรขึ้นถ้าสิ่งต่าง ๆ สามารถเกิดขึ้นได้โดยอัตโนมัติ? นี่คือ Python สิ่งต่าง ๆ เกิดขึ้นโดยอัตโนมัติใน Python
คุณจะกำหนด 'โดยอัตโนมัติ' อย่างไร?
ลองเข้าสู่บทบาทของนักพัฒนาอัลกอริทึม สิ่งต่าง ๆ ชัดเจน: ฉันไม่ต้องการสนใจประสิทธิภาพการทำงานนอกเหนือจากอัลกอริทึมหลักของฉัน เว้นแต่ฉันจะต้องทำจริง ๆ ฉันไม่อยากรู้เรื่องการสลับโมเดลเข้าและออก
แน่นอนว่าเราไม่สามารถทำเช่นนั้นได้ แต่เราสามารถพยายามลดการรบกวนที่เราต้องแนะนำให้กับโค้ดอัลกอริทึมได้ สิ่งนี้ทำให้ฉันนึกถึงการแก้ไขลิงก์ของไลบรารี gevent มันแก้ไข (หลัก ๆ) ไลบรารี socket แทนที่ด้วย gevent.socket ซึ่งสามารถสลับไปยัง greenlets อื่น ๆ เมื่อ IO จะบล็อก คล้ายกับ goroutine (จริง ๆ แล้ว gevent เก่ากว่า Golang!)
เนื่องจากเราใช้ไลบรารี HuggingFace (transformers, diffusers) ในการโหลดโมเดลในขณะนั้น เป้าหมายจึงชัดเจน: เราแนะนำการเรียก monkey-patch เพียงครั้งเดียว และโค้ดที่เหลือควรยังคงไม่เปลี่ยนแปลง XXXPipeline.from_pretrained(...) ควรจะเร็วขึ้นมาก
ข้อเท็จจริงบางประการ การตัดสินใจที่ชัดเจน และสมมติฐาน
Overmind เป็นไลบรารีแคช มันแคชผลการเรียกโหลดโมเดลลงในหน่วยความจำของระบบและสร้างมันขึ้นมาใหม่อย่างรวดเร็วในภายหลัง
เราข้ามการพูดถึงวิธีการที่ monkey-patching ถูกนำไปใช้ นั่นเป็นรายละเอียดที่ไม่น่าสนใจนัก สิ่งที่เราต้องรู้คือ มันเปลี่ยนเส้นทางการเรียก XXXPipeline.from_pretrained(...) ทั้งหมดไปที่ overmind.api.load(XXXPipeline.from_pretrained, ...)
เราใช้ pickle เพื่อจัดเก็บผลลัพธ์แคชของเราเนื่องจาก... เราไม่มีทางเลือก และ torch.save เองก็ใช้ pickle มันแปลกที่จะไม่ใช้มัน
เราใช้สถาปัตยกรรม client/server เนื่องจากเราไม่ต้องการทำให้แคชของเราเป็นโมฆะเมื่อกระบวนการสิ้นสุด มีการเรียก subprocess หลายครั้งที่อาจได้รับประโยชน์จากมัน
เราสมมติว่าพารามิเตอร์ XXXPipeline.from_pretrained เป็นสิ่งที่สามารถแฮชได้ง่าย (str และสิ่งที่คล้ายกัน) และโมเดลอื่น ๆ ที่โหลดโดย overmind (อธิบายภายหลัง)
ชื่อ overmind ยืมมาจาก Starcraft ตามที่คุณอาจเดาได้
สร้างมันขึ้นมาใหม่อย่างรวดเร็ว!
เราไม่สามารถบันทึกผลลัพธ์ pickle.loads ในหน่วยความจำและเรียกมันว่าจบได้ ในสถานการณ์ที่เตรียมพร้อมแล้ว แคชหน้าเพจของ Linux ทำงานของมันในการแคชโมเดลบนดิสก์และเรายังคงเห็นเวลาในการโหลดที่วัดได้ในระดับหลายสิบวินาที
ความไม่มีประสิทธิภาพมาจากการคัดลอกหน่วยความจำ ใน Python แม้แต่การสร้างวัตถุหลายล้านชิ้นก็จะใช้เวลาไม่เกินหลายร้อยมิลลิวินาที อย่างไรก็ตาม สำหรับการคัดลอกหน่วยความจำขนาด 10GiB จะใช้เวลาครึ่งวินาที เราต้องหลีกเลี่ยงการคัดลอกหน่วยความจำให้มากที่สุด
โชคดีที่ส่วนใหญ่ของชิ้นส่วนหน่วยความจำขนาดใหญ่เป็น Torch tensors เราสามารถระบุเฉพาะพวกมันและละเลยส่วนที่เหลือได้อย่างปลอดภัย
จริง ๆ แล้ว ฉันได้รับความรู้เกี่ยวกับโครงสร้างภายในของ Torch tensor ในโค้ดการลดขณะทำการวิจัยกลไกการแชร์ tensor:
# คัดลอกจาก 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 ได้อย่างไร? เราสามารถจัดการหน่วยความจำของเทนเซอร์เหล่านี้ด้วยตัวเองและสร้าง UntypedStorage โดยชี้ไปที่หน่วยความจำที่เราจัดการได้หรือไม่?
คำตอบคือใช่!
เมื่อเราดูโค้ด C++ ที่ UntypedStorage ถูกสร้างขึ้น เราสามารถพบโค้ดสั้นๆ เช่นนี้ได้อย่างง่ายดาย:
// คัดลอกจาก 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 เป็น 2 สิ่งหลักที่รองรับมัน และพวกมันคือสิ่งที่เราสนใจ
สุดท้ายนี้ เรารู้ว่าเราควรทำอะไร: สร้างฟังก์ชันที่รับวัตถุ memoryview และเปลี่ยนมันเป็น UntypedStorage โดยไม่ต้องคัดลอก ด้วยความสามารถในการสร้าง UntypedStorage จาก memoryview ข้อมูลเทนเซอร์จริงไม่จำเป็นต้องอยู่ในสตรีม pickle ลดขนาดข้อมูลที่เราต้องคัดลอกไปมาอย่างมาก
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 เต็มไปด้วย system calls ที่น่าสนใจ memfd_create เป็นสิ่งที่เราสนใจ: มันให้ fd ที่แสดงถึงการจัดสรรหน่วยความจำที่ไม่ระบุชื่อ คุณสามารถทำการดำเนินการไฟล์ทุกประเภทกับมัน: อ่าน, เขียน, และแน่นอน mmap หากเราสามารถแชร์ fd ได้ เราก็สามารถแชร์หน่วยความจำได้
การแชร์ fd มีวิธี 'มาตรฐาน' แต่ลึกลับในการทำ: sendmsg กับ SCM_RIGHTS เราสามารถใช้ประโยชน์จากไลบรารีเพื่อช่วยเราซ่อนรายละเอียดที่น่ากลัวของกระบวนการ sendmsg แต่เรายังคงต้องทำการประสานงานระหว่างกระบวนการเซิร์ฟเวอร์และไคลเอนต์ เราตัดสินใจใช้วิธีแฮ็กที่นี่: เพียงแค่เปิด /proc/{pidof(server)}/fd/{memfd} ที่ฝั่งไคลเอนต์ ในขณะที่ไม่เคยปิด fd ที่ฝั่งเซิร์ฟเวอร์ overmind การสื่อสารที่จำเป็นเพียงอย่างเดียวคือคู่ (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การรวมเข้ากับ pickling
ตามที่เราได้พูดคุยกันก่อนหน้านี้ เราจำเป็นต้องแก้ไขกระบวนการ pickling ของ UntypedStorage คล้ายกับที่ได้ดำเนินการใน torch.multiprocessing.reductions เรากำหนดฟังก์ชันลดที่กำหนดเองสำหรับ pickle:
# Hoarder และ borrower เป็น wrapper สำหรับ SharedMemory ข้างต้น ประกอบด้วย
# สิ่งที่น่าเบื่อเช่น memory arena เป็นต้น
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) # Zero-copy!
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 ถูกออกแบบมาสำหรับ 'การส่งเทนเซอร์ไปยัง subprocess' ดูเหมือนจะเหมือนกันแต่มีความแตกต่างเล็กน้อย
- PyTorch ใช้ POSIX shm เพื่อแชร์หน่วยความจำ ซึ่งอยู่ภายใต้ข้อจำกัดที่กล่าวถึงก่อนหน้านี้
- สำหรับทุกเทนเซอร์ (หรือ
UntypedStorage) PyTorch จัดสรรวัตถุ POSIX shm เฉพาะสำหรับมัน แม้ว่ามันจะมีเพียง 4 ไบต์ แต่ละวัตถุใช้ fd - PyTorch ยกเลิกการจัดสรร POSIX shm เมื่อมันถูก unpickled ทำให้ไม่เหมาะสมสำหรับความต้องการของเรา เราจำเป็นต้อง deserialize สตรีม pickle เดียวกันหลายครั้ง
- มีตรรกะการแชร์ที่เกี่ยวข้องกับ CUDA มากมาย ซึ่งเป็นเสียงรบกวนและปัญหาสำหรับกรณีการใช้งานของเรา
ทำไมคุณถึงบอกว่า 'ข้อมูลเทนเซอร์ถูกคัดลอกหลายครั้ง'?
สำหรับ torch.load บนดิสก์ทั่วไป:
- ไฟล์
torch.saveบนดิสก์ถูกอ่านเข้าสู่หน่วยความจำ - รับข้อมูล
torch.UntypedStorageจริงเป็นbytesโดยการแยกไฟล์ Zip (torch.saveสร้างไฟล์ zip) - โค้ด C++ จะคัดลอกข้อมูลไปยังหน่วยความจำที่จัดการเองในตัวสร้าง
torch.UntypedStorage
สำหรับ pickle.dumps แบบง่ายและ pickle.loads ในภายหลัง:
- สตรีม pickle ที่สร้างขึ้นจะฝังสตรีม pickle อื่นภายใน
pickle.loadsจะคัดลอกสตรีมภายในไปยังbytesใหม่ - ข้อมูล
torch.UntypedStorageฝังอยู่ในสตรีม pickle ภายใน จะมีการคัดลอกอีกครั้งเมื่อสร้างtorch.UntypedStorage - โค้ด C++ จะคัดลอกข้อมูลไปยังหน่วยความจำที่จัดการเองในตัวสร้าง
torch.UntypedStorage
diffusers มีโมดูลแบบไดนามิก
ที่เก็บโมเดลสามารถรวมไฟล์ Python ที่ถูกนำเข้าในขณะรันไทม์ไปยังเนมสเปซ diffusers_modules ลูกค้าไม่มีไฟล์เหล่านี้ใน sys.path ทำให้การ unpickling ล้มเหลว โชคดีที่ 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 เมื่อเราเริ่มต้น CUDA และ torch ในเซิร์ฟเวอร์ overmind แล้ว ไม่มีวิธีง่ายๆ ในการยกเลิกการเริ่มต้น ซึ่งอาจทำให้เกิดปัญหาสำหรับงานจริง (ส่วนใหญ่คือ VRAM ที่ใช้งานได้น้อยลง) ดังนั้นเราจึงแก้ไขเซิร์ฟเวอร์ของเราให้สร้าง subprocess โหลดเข้าไปในหน่วยความจำที่ใช้ร่วมกัน และยุติการทำงาน สิ่งนี้ช่วยปรับปรุงความเสถียรของเซิร์ฟเวอร์ overmind
พารามิเตอร์ที่ถูกควอนไทซ์เป็นคลาสย่อยพิเศษที่จัดทำโดย bitsandbytes พวกเขาไม่ได้ถูกออกแบบมาให้ 'picklability' ดังนั้นเราต้องทำเอง
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 มาพร้อมกับ hooks และ monkey-patches ที่ไม่สามารถ pickle ได้ เราต้องลบออก:
from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None) # ลบ warning monkeypatches
model.__dict__.pop('cuda', None)เรายังพบปัญหาที่ฟังก์ชันถูกซ้อนอยู่ภายในฟังก์ชันอื่น (แทนที่จะอยู่ในระดับบนสุด) ซึ่งทำให้ไม่สามารถ picklable ได้ เราพยายามหาทางแก้ไข แต่ไม่สำเร็จ เราต้องเปลี่ยน pickle ของเราจากที่มีให้ใน stdlib เป็น dill เพื่อ pickle สิ่งนี้ dill ทรงพลังมากกว่า แต่เป็นการดำเนินการ Python ล้วนๆ ซึ่งช้ากว่าเวอร์ชันไลบรารีมาตรฐาน โชคดีที่ต้นทุนนี้จะถูกจ่ายเพียงครั้งเดียวเมื่อเราโหลดโมเดลครั้งแรก (มีผลเฉพาะกับการ pickling ไม่ใช่ unpickling)
การสนับสนุน stable-fast
stable-fast สร้างผลลัพธ์ torch.compile ซึ่งไม่สามารถ pickle ได้ แต่ด้วย torch.jit.save เราสามารถบันทึกผลลัพธ์เป็นไฟล์ zip ได้ ฟังดูไม่มีประสิทธิภาพ แต่ก็ดีกว่าไม่มีอะไรเลย
ด้วยเพียง torch.jit.save ไม่เพียงพอที่จะ pickle ผลลัพธ์ stable-fast stable-fast ใช้กระบวนการ 'flatten' เพื่อทำให้โมดูล Torch สามารถติดตามได้ เมื่อพบสิ่งที่ไม่รู้จัก (เช่น คลาส dataclass) มันจะไม่ serialize แต่จะเก็บการอ้างอิงไปยังคลาสจริงเท่านั้น เราได้แก้ไขตรรกะที่เกี่ยวข้องเพื่อเก็บคลาสที่ถูก pickle ไว้ในสตรีม 'flatten'
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ก็มีปัญหาการคัดลอกหน่วยความจำเช่นกัน ดังนั้นเราจึงเขียน wrapper ง่าย ๆ เพื่อโหลดผ่านโปรโตคอลบัฟเฟอร์ของ 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')ตามที่เราได้กล่าวถึงก่อนหน้านี้ อาร์กิวเมนต์ของฟังก์ชันถูกสันนิษฐานว่าเป็นวัตถุที่ง่ายและสามารถ picklable ได้ง่าย แต่รูปแบบนี้ทำลายสมมตินั้น เพื่อจัดการกับสิ่งนี้ เราได้เพิ่มตรรกะพิเศษ: ผลลัพธ์ที่แคชแต่ละรายการจะได้รับ ID แนบ หากวัตถุนั้นถูกใช้เป็นอาร์กิวเมนต์ในการเรียกอื่น ไคลเอนต์จะแทนที่ด้วย ID ของมัน และเซิร์ฟเวอร์สามารถกู้คืนวัตถุจริงตาม ID ได้
โมเดล pipeline ที่ได้จะมีการอ้างอิงถึง vae เพื่อความเรียบง่าย เราเพียงแค่ pickle มันโดยตรงที่นี่ อย่างไรก็ตาม เมื่อย้าย UntypedStorage จริงไปยังหน่วยความจำที่ใช้ร่วมกัน เราจะลดข้อมูลที่ซ้ำกันใด ๆ
เราอาจใช้กลไก persistent_id ของ pickle แต่ฉันไม่ได้ลองเส้นทางนี้ นั่นน่าเสียดายเล็กน้อย
การวัดประสิทธิภาพ
และตอนนี้สำหรับส่วนที่ทุกคนชอบดู
เราใช้สคริปต์รูปแบบ VAE ของส่วนสุดท้ายเพื่อทำการทดสอบของเรา
| การทดสอบ | vae | depth | edge | pipeline | to('cuda') | รวม |
|---|---|---|---|---|---|---|
| w/o, 1st | 1.18 | 0.98 | 1.41 | 1.65 | 0.91 | 6.16 |
| w/o, 2nd | 1.15 | 0.96 | 0.97 | 1.65 | 0.89 | 5.66 |
| w/o, 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') ยังคงมีอยู่
เมื่อรวมขนาดของไฟล์โมเดลที่ถูกทำให้เป็นซีเรียลทั้งหมดแล้ว คาดว่าทั้ง pipeline จะใช้หน่วยความจำประมาณ 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 การใช้งานหน่วยความจำของระบบลดลงอย่างมาก เราไม่ได้ประสบปัญหาการขาดแคลนหน่วยความจำของระบบ แต่ถ้าเราเคยประสบ นี่จะเป็นชัยชนะที่ยิ่งใหญ่
ต่อมา เราพบว่ามันเป็นการเพิ่มประสิทธิภาพที่ยอดเยี่ยมสำหรับนักพัฒนาอัลกอริทึมและ pipeline ของเรา สำหรับแต่ละวงจรการแก้ไข-ตรวจสอบ เราสามารถประหยัดเวลาในการโหลดได้ 10 ถึง 20 วินาที ซึ่งสามารถเพิ่มขึ้นเป็นจำนวนมากได้ ที่สำคัญกว่านั้น วินาทีที่ประหยัดได้สามารถทำให้นักพัฒนาอยู่ในโฟลว์ได้
Github
เราเปิดซอร์สมันบน Github เราจะยินดีถ้ามันช่วยได้
