限られたリソースの中でLLMを使う

Accelerate を使う

大規模なLLMをそのままの状態でメモリ上に展開すると、膨大なリソースを消費します。例えば、自作PCにGPUを2枚挿しして13BのLLMを実行する場合、モデルの~B × 2GBだけリソースが消費されると仮定すると、7Bのモデルでは約14GBのリソースが必要です。一部のゲーミング用GPUにはVRAMが16GBや24GBのものもあり、14GB程度であれば比較的簡単に利用できます。しかし、70Bなどの巨大なモデルになるとリソースが140GB以上必要なため、個人での運用はリソースとコストの面でかなり厳しくなります。

これまでの話では、1つのモデルを1つのリソース（例えばVRAM）に展開することを前提としてきましたが、展開先をVRAM、RAM、ROMなど複数に分散できる場合、個人の環境でも動かせる可能性があります。このモデルを各リソースを跨いで配置する際に使用するライブラリが「Accelerate」です。

まず最初に、Accelerateではモデルのパラメータをいくつかのファイルに分割して保存します（シャードチェックポイント）。次に、それらのパラメータをどのリソースに割り当てるかを決定します。この決定は以下の手順で行います。

1. 空の（パラメータを持たない）モデルを作成する。
2. 作成したモデルの各レイヤーごとにパラメータを配置するリソースを決定する（VRAM、RAM、ROMなど）。
   このリソースの割り当てを「device_map」と呼びます。

さらに、ここで決定されたリソースに対して、以下の手順でパラメータのロードと解放を繰り返します。

1. パラメータを割り当てたリソースにロードする。
2. ロードしたパラメータを空のモデル内に配置する。
3. 推論のためにデバイス上の重みを移動する（ROM→GPUなど）。

これにより、消費される最大リソースはシャードチェックポイントの最大容量と同等までに抑えることができます。さらに詳細な説明については、How 🤗 Accelerate runs very large models thanks to PyTorchで説明されています。

具体的なコード

具体的なコードは以下のようになります。

from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

def model_load(repo_id: str):
  config = AutoConfig.from_pretrained(repo_id)
  # 手順1. 空の（パラメータを持たない）モデルを作成する。
  with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

  # ValueError: weight is on the meta device, we need a value to put in on cpu.
  # Ref: https://github.com/huggingface/accelerate/issues/1289
  model.tie_weights()

  # 手順2. 作成したモデルのレイヤごとのパラメータを配置するリソースを決める
  max_memory = {0: "10GiB", "cpu": "8GiB"}
  device_map = infer_auto_device_map(model, max_memory=max_memory)
  print(device_map)

  # 手順3. パラメータを決められたリソース上にロードする。
  model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map = device_map,
    config=config
  )

if __name__ == '__main__':
  repo_id = 'elyza/ELYZA-japanese-Llama-2-7b-fast-instruct'
  model_load(repo_id)

ここでは、手順2のdevice_map作成をinfer_auto_device_mapを使用して行っていますが、これは使用するリソースに上限 (max_memory) を設けることができるためです。max_memoryはリソース名をキー、リソース上限を値に持つ辞書型として与えられます。上記の場合、0は1台目のGPUを、cpuはRAMを表しています。またSSDやHDDもリソースとして使用する場合には、リソース名にdiskとして設定します。

device_mapが正常に作成されると以下のような出力が得られます。model.layers.x (x: 0を含む整数)はモデルの層を、0やcpuはパラメータが配置されるリソースを表しています。この出力はモデルの構造によって異なります。

OrderedDict([('model.embed_tokens', 0),
             ('model.embed_dropout', 0),
             ('model.layers.0', 0),
             ('model.layers.1', 0),
             ('model.layers.2', 0),
             ('model.layers.3', 0),
             ('model.layers.4', 0),
             ('model.layers.5', 0),
             ('model.layers.6', 0),
             ('model.layers.7', 0),
             ('model.layers.8', 0),
             ('model.layers.9', 0),
             ('model.layers.10', 0),
             ('model.layers.11', 0),
             ('model.layers.12', 0),
             ('model.layers.13', 0),
             ('model.layers.14', 0),
             ('model.layers.15', 0),
             ('model.layers.16', 0),
             ('model.layers.17', 0),
             ('model.layers.18', 0),
             ('model.layers.19', 0),
             ('model.layers.20', 0),
             ('model.layers.21.self_attn', 0),
             ('model.layers.21.mlp.gate_up_proj', 0),
             ('model.layers.21.mlp.down_proj', 'cpu'),
             ('model.layers.21.mlp.activation_fn', 'cpu'),
             ('model.layers.21.input_layernorm', 'cpu'),
             ('model.layers.21.resid_attn_dropout', 'cpu'),
             ('model.layers.21.resid_mlp_dropout', 'cpu'),
             ('model.layers.21.post_attention_layernorm', 'cpu'),
             ('model.layers.22', 'cpu'),
             ('model.layers.23', 'cpu'),
             ('model.layers.24', 'cpu'),
             ('model.layers.25', 'cpu'),
             ('model.layers.26', 'cpu'),
             ('model.layers.27', 'cpu'),
             ('model.layers.28', 'cpu'),
             ('model.layers.29', 'cpu'),
             ('model.layers.30', 'cpu'),
             ('model.layers.31', 'cpu'),
             ('model.norm', 'cpu'),
             ('lm_head', 'cpu')])

infer_auto_device_mapでは、モデルのdtype（パラメータを何ビットで表現するかの設定値）を指定することができますが、値を指定するとdevice_mapが壊れる可能性があるため注意が必要です（私はこのことに気づくまで3時間かかりました...）。device_mapが破損した場合、以下のような出力が得られます。

OrderedDict([('', 0)])

最後に、リソースごとに上限を設定しない場合、infer_auto_device_mapは不要であり、AutoModelForCausalLM.from_pretrainedのdevice_mapをautoに設定するだけで自動的にリソースの割り当てを行います。

  # 手順2. 作成したモデルのレイヤごとのパラメータを配置するリソースを決める
  # max_memory = {0: "10GiB", "cpu": "8GiB"}
  #device_map = infer_auto_device_map(model, max_memory=max_memory, dtype="float16")

  # 手順3. パラメータを決められたリソース上にロードする。
  model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map = "auto",
    config=config
  )