Tóm tắt Việc tải mô hình ML rất chậm, ngay cả khi đã làm nóng bộ nhớ đệm trang của Linux. Vì vậy, chúng tôi đã xây dựng một thư viện để làm cho nó nhanh hơn. Có một số chi tiết kỹ thuật thú vị mà chúng tôi muốn chia sẻ, vì vậy chúng tôi đã viết blog này. Thư viện cũng tạo ra một tác động bất ngờ, được thảo luận ở phần kết.
Lý do
Mọi chuyện bắt đầu từ 2 năm trước, khi chúng tôi triển khai thử nghiệm đầu tiên của chế độ tạo lowpoly. Chế độ lowpoly không diễn ra tốt, nó tạo ra kết quả kém từ quan điểm ngày nay, nhưng chúng tôi đã phải trả giá rất nhiều cho nó -- một GPU chuyên dụng chỉ xử lý được số lượng tác vụ một chữ số mỗi ngày. Nó có các trọng số tinh chỉnh, đủ lớn để đẩy tất cả các trọng số mô hình khác ra khỏi VRAM. Tệ hơn, chúng tôi có thể có 3 mô hình như vậy (không nhớ chính xác số lượng), chúng chiếm một phần đáng kể trong hạ tầng suy luận của chúng tôi, tạo ra một tỷ lệ hiệu quả khá khắc nghiệt. Và không, chúng tôi không thể tải mô hình một cách ngây thơ chỉ khi cần, nó tốn 30 giây, lớn hơn thời gian xử lý thực tế.
Chúng tôi không có kỹ sư đường ống chuyên dụng khi đó, các nhà phát triển thuật toán của chúng tôi đã cố gắng hết sức để khắc phục điều này. Vài ngày sau, mã nguồn của chúng tôi đầy rẫy các lệnh this.to('cpu') và that.to('cuda'). Cách tiếp cận này hoạt động trong một thời gian, nhưng thỉnh thoảng làm gián đoạn luồng công việc của các nhà phát triển thuật toán của chúng tôi. Sẽ ra sao nếu mọi thứ có thể xảy ra một cách tự động? Đây là Python, mọi thứ thực sự có thể xảy ra tự động trong Python.
Bạn định nghĩa 'tự động' như thế nào?
Hãy nhập vai một nhà phát triển thuật toán. Mọi thứ khá rõ ràng: Tôi không muốn quan tâm đến hiệu suất thời gian chạy bên ngoài thuật toán cốt lõi của mình trừ khi tôi thực sự phải làm vậy. Tôi thà không biết gì về việc hoán đổi mô hình vào và ra.
Tất nhiên chúng tôi không thể đạt được điều đó, nhưng chúng tôi có thể cố gắng giảm thiểu sự xâm nhập mà chúng tôi phải giới thiệu vào mã thuật toán. Điều này làm tôi nhớ đến việc vá lỗi của thư viện gevent, nó vá (chủ yếu) thư viện socket, thay thế nó bằng gevent.socket có thể chuyển sang các greenlet khác khi IO sẽ bị chặn, giống như một goroutine (thực ra gevent còn cũ hơn Golang!).
Vì chúng tôi chỉ sử dụng các thư viện của HuggingFace (transformers, diffusers) để tải mô hình vào thời điểm đó, mục tiêu trở nên rõ ràng: Chúng tôi chỉ giới thiệu một cuộc gọi vá lỗi, và phần còn lại của mã nên không thay đổi, XXXPipeline.from_pretrained(...) nên nhanh hơn nhiều.
Một số Sự thật, Quyết định Rõ ràng và Giả định
Overmind là một thư viện bộ nhớ đệm, nó lưu trữ kết quả của các cuộc gọi tải mô hình vào bộ nhớ hệ thống và sau đó tái tạo nó nhanh chóng.
Chúng tôi bỏ qua việc thảo luận về cách thực hiện vá lỗi, đó là một chi tiết không quá thú vị. Tất cả những gì chúng ta cần biết là, nó chuyển hướng tất cả các cuộc gọi XXXPipeline.from_pretrained(...) đến overmind.api.load(XXXPipeline.from_pretrained, ...).
Chúng tôi sử dụng pickle để tuần tự hóa kết quả bộ nhớ đệm của mình vì... chúng tôi không có lựa chọn nào khác, và bản thân torch.save cũng sử dụng pickle, thật kỳ lạ nếu không sử dụng nó.
Chúng tôi sử dụng kiến trúc client/server vì không muốn làm mất hiệu lực bộ nhớ đệm khi quá trình kết thúc. Có nhiều cuộc gọi subprocess có thể hưởng lợi từ điều này.
Chúng tôi giả định các tham số XXXPipeline.from_pretrained là những thứ có thể băm đơn giản (str và những thứ tương tự) và các mô hình khác được tải bởi overmind (được giải thích sau).
Tên overmind được mượn từ Starcraft, như bạn có thể đã đoán.
Tái tạo nhanh chóng!
Chúng tôi không thể ngây thơ lưu kết quả pickle.loads trong bộ nhớ và coi đó là xong. Rốt cuộc, trong một kịch bản đã làm nóng, bộ nhớ đệm trang của Linux đã làm tốt công việc của mình khi lưu trữ các mô hình trên đĩa và chúng tôi vẫn có thể thấy thời gian tải được đo bằng hàng chục giây.
Sự không hiệu quả đến từ việc sao chép bộ nhớ. Trong Python, ngay cả việc tạo hàng triệu đối tượng cũng không tốn hơn vài trăm ms. Tuy nhiên, đối với một bản sao bộ nhớ 10GiB, nó sẽ tốn nửa giây. Chúng tôi phải tránh sao chép bộ nhớ càng nhiều càng tốt.
May mắn thay, hầu hết các khối bộ nhớ lớn là các tensor của Torch, chúng tôi có thể chỉ cần xử lý chúng và bỏ qua phần còn lại.
Thực tế, tôi đã có kiến thức về cấu trúc nội bộ của một tensor Torch trong mã giảm thiểu khi nghiên cứu cơ chế chia sẻ tensor:
# Được sao chép từ torch.multiprocessing.reductions, phần lớn mã đã bị loại bỏ
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))Khá đơn giản: một tensor bao gồm kiểu của nó, metadata và lưu trữ cơ bản của nó. Ở đây storage có kiểu TypedStorage, nhưng thực tế TypedStorage chỉ là một lớp bọc đơn giản cho UntypedStorage. UntypedStorage là lớp thực sự giữ tất cả dữ liệu của tensor.
Nhiệm vụ của chúng ta trở nên cụ thể hơn: Làm thế nào để tránh sao chép UntypedStorage? Chúng ta có thể tự quản lý bộ nhớ của các tensor này và xây dựng UntypedStorage bằng cách trỏ đến bộ nhớ mà chúng ta quản lý không?
Câu trả lời là có!
Duyệt qua mã C++ nơi UntypedStorage được xây dựng, chúng ta có thể dễ dàng tìm thấy một đoạn mã như sau:
// Copied from torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
// ...omitting unrelated code...
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;
}Không chỉ có thể sử dụng một con trỏ, mà lớp at::DataPtr cũng có thể xử lý việc hủy, làm cho việc quản lý vòng đời trở nên đơn giản hơn nhiều.
Ở phía Python, một con trỏ đến một vùng bộ nhớ được biểu diễn bởi một đối tượng memoryview, các đối tượng này hỗ trợ giao thức buffer. Chúng ta có thể lấy một đối tượng memoryview từ nhiều thứ, bytes và mmap là 2 thứ chính hỗ trợ nó, và chúng cũng là những gì chúng ta quan tâm.
Cuối cùng, chúng ta biết mình nên làm gì: tạo một hàm chấp nhận một đối tượng memoryview và biến nó thành một UntypedStorage mà không cần sao chép. Với khả năng tái tạo UntypedStorage từ memoryview, dữ liệu tensor thực tế không cần phải có trong luồng pickle, giảm đáng kể kích thước dữ liệu mà chúng ta phải sao chép xung quanh.
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,
)
));
});
}Đó là khối xây dựng cốt lõi của overmind.
Chia sẻ các tensor!
Lưu ý: Đã có một cơ chế chia sẻ tensor trong PyTorch, nhưng nó không phù hợp với nhu cầu của chúng ta. Sẽ nói thêm về điều này sau.
Đầu tiên, chia sẻ bộ nhớ giữa client và server
Khi chúng ta thấy 'chia sẻ' và 'bộ nhớ' đi cùng nhau, tất cả chúng ta đều có xu hướng sử dụng shmget và các bạn của nó. Nó được "thiết kế" để được sử dụng như một cơ chế chia sẻ bộ nhớ, đúng không? Nhưng nó có 2 nhược điểm chính:
- POSIX shm là một tài nguyên khan hiếm, những gì bạn có thể sử dụng được xác định bởi cách sysadmin cấu hình hệ thống. Một ví dụ cực đoan nhưng phổ biến là các container Docker, mặc định bạn chỉ có 64MiB POSIX shm có thể sử dụng.
- POSIX shm tồn tại lâu hơn quá trình của bạn, bạn phải tự quản lý nó. Nếu quá trình quản lý bị buộc phải dừng, hoặc không xử lý cẩn thận, đối tượng shm có thể bị để lại trên hệ thống vô thời hạn.
Nếu bạn nhìn kỹ, Linux đầy những lời gọi hệ thống thú vị. memfd_create là một trong những cái chúng ta quan tâm: Nó cung cấp cho bạn một fd đại diện cho một phân bổ bộ nhớ ẩn danh. Bạn có thể thực hiện tất cả các thao tác tệp trên nó: đọc, ghi, và tất nhiên, mmap. Nếu chúng ta có thể chia sẻ fd, chúng ta có thể chia sẻ bộ nhớ.
Chia sẻ một fd có một cách 'chuẩn' nhưng bí ẩn để thực hiện: sendmsg với SCM_RIGHTS. Chúng ta có thể tận dụng các thư viện để giúp chúng ta ẩn đi các chi tiết đáng sợ của quá trình sendmsg, nhưng chúng ta vẫn phải thực hiện sự phối hợp giữa các quá trình máy chủ và khách hàng. Chúng ta quyết định sử dụng một mẹo ở đây: Chỉ cần mở /proc/{pidof(server)}/fd/{memfd} ở phía khách hàng, trong khi không bao giờ đóng fd ở phía máy chủ overmind. Giao tiếp duy nhất cần thiết là một bộ (pid, fd). Nó hoạt động hoàn hảo trong trường hợp của chúng ta.
Những lời trên được tóm gọn thành các dòng sau:
class SharedMemory:
@classmethod
def create(cls, shift):
# Được gọi ở phía máy chủ
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):
# Được gọi ở phía khách hàng
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):
# Được gọi ở cả hai phía
self._mmap = mmap.mmap(self._fd, size)
self._buf = memoryview(self._mmap)
return self._bufTích hợp với pickling
Như chúng ta đã thảo luận trước đó, chúng ta cần sửa đổi quá trình pickling của UntypedStorage. Tương tự như những gì đã được triển khai trong torch.multiprocessing.reductions, chúng ta định nghĩa các hàm giảm tùy chỉnh cho pickle:
# Hoarder và borrower là một lớp bọc cho SharedMemory ở trên, chứa
# những thứ nhàm chán như vùng nhớ, v.v.
def _reduce_storage(storage):
# Được gọi bởi máy chủ
device = storage.device
storage = storage.cpu()
# Lưu nội dung trong bộ nhớ chia sẻ
# `frag` chứa thông tin hoàn chỉnh cần thiết để định vị nội dung.
frag = hoarder.put(storage)
return (_rebuild_storage_on_client, (frag, device))
def _rebuild_storage_on_client(frag, device):
# Được gọi bởi khách hàng
mv = borrower.borrow(frag) # Lấy một memoryview từ bộ nhớ chia sẻ
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)Bây giờ, OvermindPickler.dumps và OvermindPickler.loads đơn giản sẽ sử dụng bộ nhớ chia sẻ để tăng tốc. Bạn có thể dừng đọc ở đây nếu bạn đã chán ngấy. Phần còn lại là chi tiết.
Những chi tiết quan trọng
Tại sao không sử dụng chia sẻ tensor nội bộ của PyTorch?
Đối với 'chia sẻ tensor nội bộ', tôi có ý nói torch.multiprocessing.reductions.
- Ở cấp độ cao, phương pháp của PyTorch được thiết kế để 'truyền tensor đến subprocess', có vẻ giống nhau nhưng có sự khác biệt tinh tế.
- PyTorch sử dụng POSIX shm để chia sẻ bộ nhớ, chịu giới hạn đã đề cập trước đó.
- Đối với mỗi tensor (hoặc
UntypedStorage), PyTorch phân bổ một đối tượng POSIX shm riêng biệt cho nó, ngay cả khi nó chỉ chứa 4 byte. Mỗi đối tượng tiêu tốn một fd. - PyTorch giải phóng POSIX shm ngay khi chúng được unpickled, làm cho nó không phù hợp với nhu cầu của chúng ta. Chúng ta cần giải mã cùng một luồng pickle nhiều lần.
- Có rất nhiều logic chia sẻ liên quan đến CUDA, điều này là tiếng ồn và rắc rối thuần túy cho trường hợp sử dụng của chúng ta.
Tại sao bạn nói 'dữ liệu tensor được sao chép nhiều lần'?
Đối với một torch.load điển hình trên đĩa:
- Tệp
torch.savetrên đĩa được đọc vào bộ nhớ. - Lấy dữ liệu thực tế của
torch.UntypedStoragedưới dạngbytesbằng cách giải nén tệp Zip (hàmtorch.savetạo ra một tệp zip). - Mã C++ sẽ sao chép dữ liệu vào bộ nhớ được quản lý của nó trong hàm khởi tạo
torch.UntypedStorage.
Đối với một pickle.dumps đơn giản và sau đó là pickle.loads:
- Luồng pickle được tạo ra nội bộ nhúng một luồng pickle khác,
pickle.loadssẽ sao chép luồng bên trong vào mộtbytesmới. - Dữ liệu
torch.UntypedStorageđược nhúng trong luồng pickle bên trong, một bản sao khác xảy ra khi khởi tạotorch.UntypedStorage. - Mã C++ sẽ sao chép dữ liệu vào bộ nhớ được quản lý của nó trong hàm khởi tạo
torch.UntypedStorage.
diffusers có một mô-đun động
Các kho mô hình có thể bao gồm các tệp Python được nhập vào thời gian chạy vào không gian tên diffusers_modules. Khách hàng không có chúng trong sys.path, gây ra lỗi khi giải nén. May mắn thay, diffusers sẽ ghi các tệp Python động này lên đĩa, vì vậy chúng ta chỉ cần nhập mô-đun và hoàn thành công việc.
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)Hỗ trợ cho bitsandbytes
Điều phiền phức nhất về việc hỗ trợ bitsandbytes là quá trình lượng tử hóa xảy ra trên GPU. Một khi chúng ta đã khởi tạo CUDA và torch trong máy chủ overmind, không có cách dễ dàng để hủy khởi tạo nó, điều này có thể gây ra vấn đề cho các khối lượng công việc thực tế (chủ yếu là VRAM ít sử dụng hơn). Do đó, chúng tôi đã sửa đổi máy chủ của mình để sinh ra một quy trình con, tải nó vào bộ nhớ chia sẻ và kết thúc. Điều này cải thiện sự ổn định của máy chủ overmind.
Các tham số lượng tử hóa là các lớp con đặc biệt do bitsandbytes cung cấp. Chúng không được thiết kế với 'khả năng picklability' trong tâm trí, vì vậy chúng ta phải tự làm điều đó.
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)Các mô hình lượng tử hóa qua bitsandbytes đi kèm với các hooks và monkey-patches không thể pickle, chúng ta phải loại bỏ chúng:
from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None) # Loại bỏ cảnh báo monkeypatches
model.__dict__.pop('cuda', None)Chúng tôi cũng đã gặp phải vấn đề khi các hàm được lồng trong các hàm khác (thay vì ở cấp cao nhất), điều này làm cho chúng không thể picklable. Chúng tôi đã cố gắng giải quyết vấn đề này, nhưng không thành công. Chúng tôi phải chuyển pickle của mình từ thư viện chuẩn sang dill để pickle điều này. dill mạnh mẽ hơn nhiều, nhưng nó là một triển khai thuần Python, chậm hơn nhiều so với phiên bản thư viện chuẩn. May mắn thay, chi phí này chỉ phải trả một lần khi chúng tôi tải mô hình lần đầu tiên (chỉ ảnh hưởng đến pickling, không ảnh hưởng đến unpickling).
Hỗ trợ cho stable-fast
stable-fast tạo ra kết quả torch.compile, không thể pickle. Nhưng với torch.jit.save, chúng ta có thể lưu kết quả dưới dạng tệp zip. Điều này nghe có vẻ không hiệu quả, nhưng vẫn tốt hơn là không có gì.
Chỉ với torch.jit.save là không đủ để pickle kết quả stable-fast. stable-fast sử dụng một quá trình 'flatten' để làm cho mô-đun Torch có thể truy vết. Khi gặp phải thứ gì đó mà nó không nhận ra (ví dụ, lớp của dataclass), nó sẽ không tuần tự hóa nó, mà chỉ giữ một tham chiếu đến lớp thực tế. Chúng tôi đã vá logic liên quan để thực sự lưu trữ một lớp đã được pickle trong luồng '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_dataclassCó hai mẹo nữa ở đây:
- Chúng tôi đóng gói lại tệp ZIP với
ZIP_STORED, vì vậy chúng tôi không cần giải nén tệp ZIP cho mỗi lần tải sau đó. - Giao diện
torch.jit.loadcũng gặp vấn đề sao chép bộ nhớ, vì vậy chúng tôi đã viết một trình bao bọc đơn giản để tải nó qua giao thức bộ đệm Python, giống như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); // Không sao chép!
return import_ir_module(std::move(cu), in, ...);
}
);
}Mẫu vae=vae
Cơ sở mã của chúng tôi có một cái gì đó như thế này, nó cố gắng tải một mô hình với một mô hình đã tải trước đó làm đối số của nó:
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, # Ở đây!
controlnet=[controlnet_edge, controlnet_depth], # và Ở đây!
torch_dtype=torch.float16,
safety_checker=None,
)
pipeline.to('cuda')Như chúng tôi đã đề cập trước đó, các đối số hàm được giả định là các đối tượng đơn giản, dễ dàng có thể được picklable, nhưng mẫu này phá vỡ giả định đó. Để xử lý điều này, chúng tôi đã thêm logic đặc biệt: mỗi kết quả được lưu trữ có một ID đính kèm. Nếu đối tượng đó được sử dụng như một đối số trong một cuộc gọi khác, máy khách sẽ thay thế nó bằng ID của nó và máy chủ sau đó có thể khôi phục đối tượng thực dựa trên ID.
Mô hình pipeline kết quả sẽ chứa một tham chiếu đến vae. Để đơn giản, chúng tôi chỉ cần pickle nó trực tiếp ở đây. Tuy nhiên, khi di chuyển UntypedStorage thực tế vào bộ nhớ chia sẻ, chúng tôi loại bỏ bất kỳ dữ liệu lặp lại nào.
Chúng tôi có thể đã sử dụng cơ chế persistent_id của pickle, nhưng tôi đã không thử theo hướng này. Đó là một chút đáng tiếc.
Đánh giá hiệu năng
Và bây giờ là phần mà mọi người đều thích xem.
Chúng tôi sử dụng kịch bản mẫu VAE của phần cuối cùng để thực hiện kiểm tra của chúng tôi.
| Kiểm tra | vae | depth | edge | pipeline | to('cuda') | Tổng cộng |
|---|---|---|---|---|---|---|
| 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 |
Như bạn có thể thấy, lần tải đầu tiên với overmind mất 24.2 giây, lâu hơn đáng kể so với khi không sử dụng nó. Tuy nhiên, trong các lần tải tiếp theo, chỉ còn chi phí .to('cuda').
Cộng tổng kích thước của tất cả các tệp mô hình được tuần tự hóa, toàn bộ pipeline ước tính sử dụng khoảng 5808 megabyte bộ nhớ. Một thử nghiệm nhanh cho kết quả tương tự.
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)Đã thử nghiệm trên Intel i9-11900K + GeForce RTX 4090.
Tác Dụng Phụ Không Mong Đợi (Tích Cực!)
Động lực chính của chúng tôi khi xây dựng overmind là để cho phép chuyển đổi nhanh chóng trọng số mô hình trong quá trình suy luận. Trong khi nó đã phục vụ mục đích của mình, chúng tôi đã phát hiện ra một số lợi ích bổ sung trong quá trình.
Chúng tôi triển khai nhiều phiên bản của ứng dụng, mỗi phiên bản cho một GPU. Do đó, sẽ có 8 tiến trình trên mỗi node. Sau khi chúng tôi triển khai overmind, việc sử dụng bộ nhớ hệ thống đã giảm đáng kể. Chúng tôi không gặp phải tình trạng thiếu bộ nhớ hệ thống, nhưng nếu có, điều này sẽ là một chiến thắng lớn.
Sau đó, chúng tôi nhận thấy nó là một sự thúc đẩy lớn cho các nhà phát triển thuật toán và pipeline của chúng tôi. Đối với mỗi vòng lặp sửa đổi-xác minh, chúng tôi có thể tiết kiệm 10 đến 20 giây thời gian tải, điều này có thể cộng lại thành một con số lớn. Quan trọng hơn, những giây tiết kiệm được có thể giữ cho các nhà phát triển trong trạng thái làm việc hiệu quả.
Github
Chúng tôi đang mở mã nguồn trên Github, chúng tôi sẽ rất vui nếu nó giúp ích.
