WSL2でOpenELMを試してみる

Appleが公開した「オープンソースの効率的言語モデルのファミリー」らしいOpenELM。

apple/OpenELM · Hugging Face

OpenELMは、レイヤーごとのスケーリング戦略を使用して、トランスフォーマーモデルの各レイヤー内でパラメーターを効率的に割り当て、精度の向上につながります。CoreNetライブラリを使用してOpenELMモデルを事前トレーニングしました。270M、450M、1.1B、および3Bパラメーターを備えた事前トレーニング済みモデルと命令調整済みモデルの両方をリリースします。

私たちの事前トレーニングデータセットには、RefinedWeb、重複排除されたPILE、RedPajamaのサブセット、およびDolma v1.6のサブセットが含まれており、合計約1.8兆トークンになります。これらのデータセットを使用する前に、ライセンス契約と条件を確認してください。

Model cardの冒頭部をGoogle翻訳

今回はInstructモデルの以下の4つを試してみます。

• OpenELM-270M-Instruct

• OpenELM-450M-Instruct

• OpenELM-1_1B-Instruct

• OpenELM-3B-Instruct

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

venv構築

python3 -m venv openelm
cd $_
source bin/activate

パッケージのインストール

pip install torch transformers accelerate

2. 流し込むコード

generate_openelm.pyをみるに、tokenizerとして

tokenizer: Union[str, AutoTokenizer] = 'meta-llama/Llama-2-7b-hf',

が指定されていました。モデルとは別にtokenizerの指定が必要となるため、いつものscriptにtokenizer-pathが指定できるように処理を追加しました。こちらをquery.pyとして保存します。

import sys
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
from typing import List, Dict
import time

# argv
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--tokenizer-path", type=str, default=None)
parser.add_argument("--no-chat", action='store_true')
parser.add_argument("--no-use-system-prompt", action='store_true')
parser.add_argument("--max-tokens", type=int, default=256)

args = parser.parse_args(sys.argv[1:])

model_id = args.model_path
if model_id == None:
    exit

is_chat = not args.no_chat
use_system_prompt = not args.no_use_system_prompt
max_new_tokens = args.max_tokens

tokenizer_id = model_id
if args.tokenizer_path:
    tokenizer_id = args.tokenizer_path

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    tokenizer_id,
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    #torch_dtype=torch.bfloat16,
    device_map="auto",
    #device_map="cuda",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
#if torch.cuda.is_available():
#    model = model.to("cuda")

streamer = TextStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True
)

DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"


def q(
    user_query: str,
    history: List[Dict[str, str]]=None
) -> List[Dict[str, str]]:
    # generation params
    generation_params = {
        "do_sample": True,
        "temperature": 0.8,
        "top_p": 0.95,
        "top_k": 40,
        "max_new_tokens": max_new_tokens,
        "repetition_penalty": 1.1,
    }
    #
    start = time.process_time()
    # messages
    messages = ""
    if is_chat:
        messages = []
        if use_system_prompt:
            messages = [
                {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
            ]
        user_messages = [
            {"role": "user", "content": user_query}
        ]
    else:
        user_messages = user_query
    if history:
        user_messages = history + user_messages
    messages += user_messages
    # generation prompts
    if is_chat:
        prompt = tokenizer.apply_chat_template(
            conversation=messages,
            add_generation_prompt=True,
            tokenize=False
        )
    else:
        prompt = messages
    input_ids = tokenizer.encode(
        prompt,
        add_special_tokens=True,
        return_tensors="pt"
    )
    print("--- prompt")
    print(prompt)
    print("--- output")
    # 推論
    output_ids = model.generate(
        input_ids.to(model.device),
        streamer=streamer,
        **generation_params
    )
    output = tokenizer.decode(
        output_ids[0][input_ids.size(1) :],
        skip_special_tokens=True
    )
    if is_chat:
        user_messages.append(
            {"role": "assistant", "content": output}
        )
    else:
        user_messages += output
    end = time.process_time()
    ##
    input_tokens = len(input_ids[0])
    output_tokens = len(output_ids[0][input_ids.size(1) :])
    total_time = end - start
    tps = output_tokens / total_time
    print(f"prompt tokens = {input_tokens:.7g}")
    print(f"output tokens = {output_tokens:.7g} ({tps:f} [tps])")
    print(f"   total time = {total_time:f} [s]")
    return user_messages

print('history = ""')
print('history = q("ドラえもんとはなにか")')
print('history = q("続きを教えてください", history)')

3. 試してみる

実行コマンド

model-pathは試したいモデルに応じて変更します。tokenizer-pathはmeta-llama/Llama-2-7b-hfを指定します。
※tokeninzerは、elyza/ELYZA-japanese-Llama-2-7b-instructでも動きました。

# 270M
CUDA_VISIBLE_DEVICES=0 python -i /path/to/query.py --model-path apple/OpenELM-270M-Instruct --tokenizer-path meta-llama/Llama-2-7b-hf --no-chat

# 450M
CUDA_VISIBLE_DEVICES=0 python -i /path/to/query.py --model-path apple/OpenELM-450M-Instruct --tokenizer-path meta-llama/Llama-2-7b-hf --no-chat

# 1.1B
CUDA_VISIBLE_DEVICES=0 python -i /path/to/query.py --model-path apple/OpenELM-1_1B-Instruct --tokenizer-path meta-llama/Llama-2-7b-hf --no-chat

# 3B
CUDA_VISIBLE_DEVICES=0 python -i /path/to/query.py --model-path apple/OpenELM-3B-Instruct --tokenizer-path meta-llama/Llama-2-7b-hf --no-chat

以下いずれのモデルも

history = q("ドラえもんとは")

と質問した結果です。

(1) 270M

使用vramは1.1GBほど。

、エクスポートしたインターネットに続く「拡大戦」を開始する。シリーズは日本の海外軍隊(ナレパン)の第2戦隊(ヴァルティア)と第1戦隊(バラスク)を役立ちます。新たな戦闘が起きている頃に争う困難力を巡る軍隊に合意します。戦闘は自由経営であり、国民や戦争者に適使用可能です。

戦争者は、中央個国国軍防衛省(CINAD)に対して決定できます。中央個国の最高効果力を持っている米軍事部門によると、この戦闘は日本に対応し
prompt tokens = 9
output tokens = 256 (50.638146 [tps])
total time = 5.055477 [s]

apple/OpenELM-270M-Instruct

(2) 450M

使用vramは1.4GBほど。

……?

It's the most famous and well-known of all Dr. Who stories, and if you haven't watched it already, you need to.

**Synopsis:** 3 Time Lords, Dr. Who, and The Master meet a mysterious being called Dr. Orin Hardy (played by Patrick Troughton), who's determined to discover the source of an ancient cosmic force that threatens the very fabric of existence.

>**[Watch the 2005 version of the story here] (https://www.youtube.com/watch?v=TqFNwR9lkJQ)**

*This video contains spoilers for the original story.*

Dr. Orin Hardy is a Time Lord who's been sent by the Time Lord King Raxadan (played by Mark Gatiss) to investigate the power source that threatens to destroy humanity. When he meets the mysterious Dr. Who, it sparks an encounter that sets in motion the timeline-changing events that define this episode.*

A decade earlier, the Time Lords journeyed to the planet Orion, an
prompt tokens = 9
output tokens = 256 (44.701477 [tps])
total time = 5.726880 [s]

apple/OpenELM-450M-Instruct

(3) 1.1B

使用vramは2.6GBほど。

?

「いつでもドラえもんですか。私の頃もおりました。」

観客は、これまでにいろんなドラえもんを見ていたけど、今度は対象となった。

彼女はその後、そのトップで固定した。

「これがみたいなドラえもんです。それでも前の時代が魅力でした。いや、今度は触れることが出来ます。」

ところで、彼女達はその道を選択し、いきなりドラえもんに進む。

「ドラえもんが好きですよ。驚きました。何とか見つけましたね。」

観客は、現在はそのイメージを保持するだけで、もう一度ドラ
prompt tokens = 9
output tokens = 256 (30.200959 [tps])
total time = 8.476552 [s]

apple/OpenELM-1_1B-Instruct

(4) 3B

使用vramは6.4GBほど。

何で違っているのか

みなさまこんにちは。浦野サイコミックです。

今回は、ドラえもんとは何で違っているのかという記事です。

そもそもドラえもんなのですか。

ダ・ウィン・フィクションとは

筆者が生まれ育ったころに、この名前の子どもが存在しました。

このおばあちゃんは、姉さんの妹でもあった筆者の母親から、ドラえもんというのは別物ではなかったと教わっていました。

それは、ダ・ウィン・フィクションだったということです。

その筆者は、その後、フィクションジャパンとは異
prompt tokens = 9
output tokens = 256 (24.824918 [tps])
total time = 10.312219 [s]

apple/OpenELM-3B-Instruct

4. まとめ

モデルサイズ毎のVRAM使用量とtoken/secをまとめます。

270M: 1.1GB, 50.63
450M: 1.4GB, 44.70
1.1B: 2.6GB, 30.20
3B: 6.4GB, 24.82

サイズが270Mだと、推論速度も速いですし、VRAMの使用量も1.1GB程度。
スペックがそんなに高くないiPhoneでも動きそうです