GTX1060 6G にて device_map を少しいじって cyberagent/open-calm-3b を動かしてみた。

win11home、i7、32G、GTX1060 6G 環境です。

cyberagent/open-calm-1b をいじっていて cyberagent/open-calm-3b ならおそらく device_map をいじってやったら動くのではないかなとやってみたら動かせたのでせっかくなので書いておこうかと。*1

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import infer_auto_device_map

model_name="cyberagent/open-calm-3b"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(model_name)
device_map = infer_auto_device_map(model, 
                                   max_memory={0: "5GiB", "cpu": "16GiB"},
                                  )
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map=device_map,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
).eval()

inputs = tokenizer("AIによって私達の暮らしは、", return_tensors="pt").to(model.device)
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=True,
        temperature=0.7,
        pad_token_id=tokenizer.pad_token_id,
    )
    
output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(output)


>>> AIによって私達の暮らしは、とても便利なものになっています。
>>> しかし、その便利さの裏側には、様々な危険性が潜んでいることを私たちは、なかなか知る機会がありません。
>>> 私たちが、日常生活で無意識に使っている、電磁波や、重金属、農薬や化学肥料など、本来、私たちが生きている地球環境に無い物質こそが、有害である可能性があります。
>>> 私たちは

model を一回読み込んで、device_map を infer_auto_device_map で設定し直してもう一度 model を読み込むっぽいです。タスクマネージャーを見ると、GPU は3Dは11％くらい、コピーが75％くらいで稼働しているので意味があるのかないのかはよくわかりませんが、とにかく動いてはいます。

*1:cyberagent/open-calm-7b だと SSD のアクセスがとんでもないことになって、いつまでも vram に上がってこないのでビビってやめました。