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라는 이름은 스타크래프트에서 가져온 것임을 짐작하셨을 것입니다.
빠르게 재구성하세요!
우리는 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의 복사를 피할 수 있을까요? 우리가 직접 텐서 메모리를 관리하고, 우리가 관리하는 메모리를 가리키는 방식으로 UntypedStorage를 구성할 수 있을까요?
답은 예입니다!
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 객체로 표현됩니다. 이 객체들은 버퍼 프로토콜을 지원합니다. 우리는 bytes와 mmap 같은 여러 가지로부터 memoryview 객체를 얻을 수 있으며, 이것들이 우리가 관심을 가지는 주요 대상입니다.
마지막으로, 우리는 무엇을 해야 하는지 알게 되었습니다: 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과 그 친구들을 사용하고 싶은 충동을 느낍니다. 이것은 메모리 공유 메커니즘으로 "설계"된 것이죠, 그렇지 않나요? 하지만 두 가지 주요 결점이 있습니다:
- POSIX shm은 희소 자원입니다. 사용할 수 있는 것은 시스템 관리자가 시스템을 어떻게 구성하느냐에 따라 결정됩니다. 극단적이지만 어디서나 볼 수 있는 예는 Docker 컨테이너입니다. 기본적으로 64MiB의 POSIX shm만 사용할 수 있습니다.
- POSIX shm은 프로세스보다 오래 지속되므로, 직접 관리해야 합니다. 관리 프로세스가 강제로 종료되거나 신중하게 처리되지 않으면 shm 객체가 시스템에 무기한 남아 있을 수 있습니다.
자세히 살펴보면, Linux에는 흥미로운 시스템 호출이 많습니다. 우리가 관심 있는 memfd_create는 익명의 메모리 할당을 나타내는 fd를 제공합니다. 이를 통해 모든 종류의 파일 작업을 수행할 수 있습니다: 읽기, 쓰기, 그리고 물론 mmap도 가능합니다. fd를 공유할 수 있다면 메모리도 공유할 수 있습니다.
fd를 공유하는 데는 '표준적이지만 난해한' 방법이 있습니다: SCM_RIGHTS와 함께 sendmsg를 사용하는 것입니다. 우리는 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파일이 메모리로 읽힙니다. - Zip 파일 추출을 통해 실제
torch.UntypedStorage데이터를bytes로 가져옵니다 (torch.save는 zip 파일을 생성합니다). - C++ 코드는
torch.UntypedStorage생성자에서 데이터를 자체 관리 메모리로 복사합니다.
단순한 pickle.dumps와 이후 pickle.loads의 경우:
- 생성된 pickle 스트림은 내부적으로 또 다른 pickle 스트림을 포함하며,
pickle.loads는 내부 스트림을 새로운bytes로 복사합니다. torch.UntypedStorage데이터는 내부 pickle 스트림에 포함되어 있으며,torch.UntypedStorage생성 시 또 다른 복사가 발생합니다.- C++ 코드는
torch.UntypedStorage생성자에서 데이터를 자체 관리 메모리로 복사합니다.
diffusers는 동적 모듈을 포함합니다
모델 저장소는 런타임에 diffusers_modules 네임스페이스로 가져오는 Python 파일을 포함할 수 있습니다. 클라이언트는 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에서 발생한다는 것입니다. overmind 서버에서 CUDA와 torch를 초기화한 후에는 이를 쉽게 해제할 방법이 없어 실제 작업 부하에 문제를 일으킬 수 있습니다 (주로 사용 가능한 VRAM이 줄어듦). 따라서 우리는 서버를 수정하여 하위 프로세스를 생성하고, 공유 메모리에 로드하고, 종료하도록 했습니다. 이는 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를 통한 양자화된 모델은 pickle되지 않는 후크와 몽키 패치가 함께 제공되므로 이를 제거해야 합니다:
from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None) # 경고 몽키 패치 제거
model.__dict__.pop('cuda', None)또한 함수가 다른 함수 내에 중첩되어 있어 (최상위 레벨이 아닌) picklable하지 않은 문제를 겪었습니다. 이를 해결하려고 시도했지만 성공하지 못했습니다. 우리는 표준 라이브러리에서 제공하는 pickle을 dill로 교체해야 했습니다. dill은 훨씬 강력하지만 순수 Python 구현으로, 표준 라이브러리 버전보다 훨씬 느립니다. 다행히도, 이 비용은 모델을 처음 로드할 때 한 번만 지불하면 됩니다 (pickling에만 영향을 미치며 unpickling에는 영향을 미치지 않습니다).
stable-fast 지원
stable-fast는 torch.compile 결과를 생성하며, 이는 pickle될 수 없습니다. 그러나 torch.jit.save를 사용하여 결과를 zip 파일로 저장할 수 있습니다. 이는 비효율적으로 보일 수 있지만, 없는 것보다는 낫습니다.
torch.jit.save만으로는 stable-fast 결과를 pickle하기에 충분하지 않습니다. stable-fast는 Torch 모듈을 추적 가능하게 만들기 위해 'flatten' 프로세스를 사용합니다. 인식하지 못하는 것을 만나면 (예: dataclass의 클래스), 이를 직렬화하지 않고 실제 클래스에 대한 참조만 유지합니다. 우리는 'flatten'된 스트림 내에 pickle된 클래스를 실제로 저장하도록 관련 로직을 패치했습니다.
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인터페이스도 메모리 복사 문제를 일으키므로, 우리는UntypedStorage처럼 Python 버퍼 프로토콜을 통해 로드하기 위한 간단한 래퍼를 작성했습니다.
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') | 총계 |
|---|---|---|---|---|---|---|
| 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') 비용만 여전히 존재합니다.
모든 직렬화된 모델 파일의 크기를 합산하면 전체 파이프라인은 약 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에 오픈 소스로 공개하고 있으며, 도움이 된다면 기쁘겠습니다.
