transformersでLLMのcache保存場所を指定して8bit量子化で推論するためのサンプルコード

こんにちは、クロガネです。
自分のメモ用に、transformersでLLMを動かすときに、特定のディレクトにcacheを保存しつつ、bitsandbytesで8bit量子化をして推論するためのコードを置いておきます。

Cドライブが狭く、VRAMが小さい時に便利です。
今回のコードではGドライブのHDDにcacheを保存しています。

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

# tokenizerとmodelのロード
# 量子化設定
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
# キャッシュ用ディレクトリ
dir_cache = r"G:\cache"

tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen3-0.6B",
    cache_dir = dir_cache, # キャッシュの指定
    )
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", 
    device_map="auto",
    quantization_config=bnb_config, # 量子化設定の指定
    cache_dir = dir_cache, # キャッシュの指定
    )

model.eval()

# チャットログ形式のprompt用辞書
d_chats = [
    {"role": "system", "content": "/think"},
    {"role": "user", "content": "Who is the cutest in Madoka Magica?"},
]

# tokenize部分
input_ids = tokenizer.apply_chat_template(
    d_chats, 
    add_generation_prompt=True, 
    return_tensors="pt",
    ).to(model.device)

# 推論部分
with torch.no_grad():
    output_ids = model.generate(
        input_ids, 
        temperature=0.1, 
        do_sample=True, 
        top_p=0.95, 
        top_k=40, 
        max_new_tokens=256,
        eos_token_id=tokenizer.eos_token_id, 
    )

# 出力のでコード部分
# 入出力が正常に設定できているか確認するために入力トークンと特殊トークンもすべてprint()している
output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output)

以上です。