TL;DR MLモデルの読み込みは、Linuxページキャッシュをウォームアップしても遅いです。そこで、私たちはそれを高速化するライブラリを作成しました。共有したい興味深い技術的詳細がいくつかあるので、このブログを書きました。このライブラリは予期しない影響も与え、そのことについては最後に説明します。
理由
すべては2年前、私たちが初めてローポリ生成モードを出荷したときに始まりました。ローポリモードはうまくいかず、今日の視点から見ると貧弱な結果を出しましたが、それに多くの費用をかけました。専用のGPUが1日に一桁のタスクしか処理しませんでした。それは微調整された重みを持ち、他のすべてのモデルの重みをVRAMから追い出すのに十分な大きさでした。さらに悪いことに、私たちはそのようなモデルを3つ持っていたかもしれません（正確な数は覚えていません）、それらは私たちの推論インフラの重要な部分を構成し、非常に厳しい効率比を作り出しました。そして、モデルをジャストインタイムで単純に読み込むことはできません。30秒かかり、実際の処理時間よりも長いです。
当時、専用のパイプラインエンジニアはいませんでしたので、私たちのアルゴリズム開発者はこれを回避するために最善を尽くしました。数日後、私たちのコードベースはthis.to('cpu')やthat.to('cuda')で散らかっていました。このアプローチはしばらくの間機能しましたが、アルゴリズム開発者の流れを時々中断しました。もし物事が自動的に起こることができたらどうでしょうか？Pythonでは、物事は自動的に起こります。
'自動的に'をどのように定義しますか？
アルゴリズム開発者の役割に飛び込んでみましょう。状況は非常に明確です：私は絶対に必要でない限り、コアアルゴリズム以外のランタイムパフォーマンスについて気にしたくありません。モデルのスワップインとアウトについて何も知らない方がいいです。
もちろんそれを達成することはできませんが、アルゴリズムコードに導入しなければならない侵入を最小限に抑えることはできます。これは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ページキャッシュがディスク上のモデルをキャッシュする役割を果たし、読み込み時間が数十秒で測定されることがあります。
非効率性はメモリコピーから来ています。Pythonでは、数百万のオブジェクトを作成しても数百ミリ秒以上かかることはありません。しかし、10GiBのメモリコピーには半秒かかります。メモリコピーをできるだけ避けなければなりません。
幸いなことに、大きなメモリチャンクのほとんどはTorchテンソルであり、それらだけを安全にアドレスし、残りを無視することができます。
実際、テンソル共有メカニズムを研究している間に、Torchテンソルの内部構造に関する知識を削減コードで得ました：
# Copied from torch.multiprocessing.reductions, most of the code is removed
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オブジェクトで表され、これらのオブジェクトはバッファプロトコルをサポートしています。多くのものから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のためのカスタムreduce関数を定義します：
# 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.loadsは内部ストリームを新しいbytesにコピーします。 torch.UntypedStorageデータは内部ピクルストリームに埋め込まれ、torch.UntypedStorageの構築時にもう一度コピーが行われます。- C++コードは、
torch.UntypedStorageコンストラクタ内でデータを自身の管理メモリにコピーします。
diffusersには動的モジュールがあります
モデルリポジトリは、実行時にdiffusers_modules名前空間にインポートされるPythonファイルを含むことができます。クライアントはこれらを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)他の関数内にネストされている関数（トップレベルではなく）が原因でピクル化できない問題にも遭遇しました。これを回避しようとしましたが、うまくいきませんでした。標準ライブラリ提供のピクルをdillに切り替える必要がありました。dillは非常に強力ですが、純粋なPython実装であるため、標準ライブラリバージョンよりもはるかに遅いです。幸いなことに、このコストはモデルを初めてロードする際に一度だけ支払われます（ピクル化にのみ影響し、アンピクル化には影響しません）。
stable-fastのサポート
stable-fastはtorch.compile結果を生成しますが、これはピクル化できません。しかし、torch.jit.saveを使用することで、結果をzipファイルとして保存できます。これは非効率的に思えますが、何もないよりはマシです。
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ここにはさらに2つのトリックがあります：
- 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); // コピーなし！
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秒かかり、これに比べて使用しない場合のロードはかなり短いです。しかし、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ごとに1つずつデプロイしています。したがって、ノードごとに8つのプロセスがあります。overmindをデプロイした後、システムメモリの使用量が劇的に減少しました。システムメモリの不足に悩まされていたわけではありませんが、もしそうだったら、これは大きな勝利だったでしょう。
後に、アルゴリズムとパイプラインの開発者にとって大きな助けとなることがわかりました。各修正-検証ループで、10から20秒のロード時間を節約でき、これは大きな数に積み重なる可能性があります。さらに重要なのは、節約された秒数が開発者をフロー状態に保つことができることです。
Github
Githubでオープンソース化しています。役に立てば嬉しいです。
