初めに
bnbやAccelerateを使用します。
VRAMとは何か、というお話ですが、コンピューター内でのビット表現などにおいて、それらのデータをどこにマップするかというお話です。
AIは、行列の計算に特化したGPUを使用し、計算速度を高めています。
このような構想は今では普通になり、ほとんどのAIの研究でGPUの使用がされています。
しかし、このような問題を経ても新たな問題はあります。それがVRAMという本来限られているはずのメモリーがweightsなどのために膨大に使用されるという問題です。
そもそもVRAMは、テクスチャなどをGPU側のアロケートし、ボトルネックを解消するために存在します。
まあ、これ自体は簡単なお話ですが、このVRAMにAIの巨大なweightsというものをメモリーにアロケートするのは、あまり良いアイデアには思えません。
しかし、企業の多くはFP16を使用しています。
一般消費者のGPUにそれをアロケートさせるというのは悪い夢のようなものです。しかし、メディアが最新モデルをローカルで使う方法などを紹介しているため、それに応えるGPUをそのまま買うという人も出てきます。
今回は、これらの問題を考えてみます。
GPUにアロケートする。
deepseek-ai/DeepSeek-R1-0528-Qwen3-8B
というモデルをアロケートしてみます。
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
model_name = "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
base_model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="cuda"
)
device_mapをcudaにすると、cuda側にアロケートするよう命令されます。
この場合、cudaに直接すべてアロケートするよう書かれています。
当たり前ですが、いつのまにかVRAMがひっ迫され限界を迎え、その旨をCUDA側から伝えられます。
CUDA out of memory. Tried to allocate 192.00 MiB. GPU 0 has a total capacity of 6.00 GiB of which 0 bytes is free. Of the allocated memory 12.63 GiB is allocated by PyTorch, and 11.53 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
https://huggingface.co/deepseek-ai/DeepSeek-R1-0528-Qwen3-8B/tree/main
どれほどの量なのかは、実際のweightsのファイルのサイズからおおまかわかるでしょう。
おおよそ16GBです。
では、これらに対処する方法を考えます。
RAM・HHD or SSDにアロケートする。
weightsをRAM側にアロケートされます。
device_mapをcpuにすることで、ramにアロケートされます。
しかし、メモリーの場合でも、限界を迎え、パソコンがかなり重たくなってしまいました。
流石に、これは使い過ぎです。制限します。
max_memory={0: "6GiB", "cpu": "6GiB"},
これで制限をかけることができ、cpuのメモリーでも読み込めなかったものは、HHDなどに行きます。
base_model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
load_in_4bit=True,
max_memory={0: "6GiB", "cpu": "6GiB"},
llm_int8_enable_fp32_cpu_offload=True
)
これを設定することでなんとか読み込むことができました。
試しに推論してみます。
inputs = tokenizer.encode("<ROLE>あなたはAIアシスタントです。ユーザーのリクエストにこたえる義務があります。</ROLE><THINK>私の役目は", return_tensors="pt").to(model.device)
outputs = model.generate(
inputs,
max_length=150,
)
generated_text = tokenizer.decode(outputs[0])
print(generated_text)
<ROLE>あなたはAIアシスタントです。ユーザーのリクエストにこたえる義務があります。</ROLE><THINK>私の役目は、ユーザーのクエリに忠実に従い、明確で簡潔な回答を提供することです。ユーザーは「AIアシスタント」という役割を指定しているため、回答はAIアシスタNTの視点で考え、ユーザーのニーズに応じる必要があります。</THINK>
と回答してくれました。