AirLLM:405BのmodelをLocalで動かすと

去年の年末頃、比較的大きなLLMのmodelでも各層毎にGPUで計算することで少ないGPUメモリでも動かすことができるAirLLMが公開されました。

その方がLlama 3.1 405B modelを8GB VRAM以下で動作させていたので、試してみました。

Crazy Challenge: Run Llama 405B on a 8GB VRAM GPU

概要

Llama-3.1-405B

言わずもがな、Meta社の現在最高峰のオープンソース、多言語対応、128Kと長いコンテキスト長のmodelです。
動作には800GB程度のVRAMが必要らしいです。

4bit量子化

16bitから4bit量子化で必要メモリ量が1/4近くになります。

unsloth/Meta-Llama-3.1-405B-Instruct-bnb-4bit · Hugging Face

層毎に分割推論

Llama-3.1-405Bは下図の通り126層から成ります。

これを各層ごとに分割しGPUにのせることで最大VRAM使用量を~6GB程度に抑えることができます。

実施内容

0. 環境

OS：Windows
CPU：Intel(R) Core i9-13900KF
RAM：128GB
GPU：RTX 4090 (24GB)

1. 実行

Codeは下記とシンプルです。
出力に時間を要すため、max_new_tokensは10と小さくしています。

from airllm import AutoModel

model = AutoModel.from_pretrained(
    "unsloth/Meta-Llama-3.1-405B-Instruct-bnb-4bit")

input_text = ['日本の首都は？',]

input_tokens = model.tokenizer(input_text,
      return_tensors="pt", 
      return_attention_mask=False, 
      truncation=True, 
      max_length=128, 
      padding=False)

generation_output = model.generate(
      input_tokens['input_ids'].cuda(), 
      max_new_tokens=10,
      return_dict_in_generate=True)

output = model.tokenizer.decode(generation_output.sequences[0])

print(output)

回答

<|begin_of_text|>日本の首都は？（What is the capital of Japan?） 東京

2. HFのCacheの容量

405GBほどになりました。

3. 計算中のPC挙動

たしかにVRAMは非常に小さく抑えられていました。
常時ストレージとGPUが行き来するので、普段静かなSSDがアクティブなのは新鮮です。

4. 出力時間

1tokenあたり126層の計算に2min要しました。

今回は10tokenで20minかかりました。

所感

AirLLMが公開された当時、すでにMistralベースをはじめ小さく優秀なmodelがでており、あまり実用の機会はありませんでした。

ただ、やはり今回のように400Bほどの大型のmodelが手元のPCで簡単に動かせることは、2年前には想像もできませんでしたし、素直に感動します。