Gemma 2B をMacで動かしてみた

Googleが公開した Gemma 2B の Instruction Tunedモデルを Transformer で動かした時の記録です。また、1.1BのTiny Llamaとの比較を行いました。

実行環境

実行環境は以下になります。

  • OS:MacOS Sonoma
  • Python:3.10.13
  • Torch:2.2.1
  • Transformers:4.38.1

Gemmaの準備

KaggleからGemma 2B のInstruction Tuned済みモデルをダウンロードします。 Kaggleにログインし以下のモデルカードでModel Valiation からTransformersタブで2b-itを選択します。

環境の準備

最新のTransformersをインストール済みの環境を前提とします。

ダウンロードしたモデルを解凍しその中に以下のソースを追加します。

llm.ipynb

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("./")
model = AutoModelForCausalLM.from_pretrained("./")

text = ""
token_ids = tokenizer.encode(text, add_special_tokens=True, return_tensors="pt")
token_length = len(token_ids[0])

outputs = model.generate(
        token_ids.to(model.device),
        max_new_tokens=token_length + 100,
        min_new_tokens=token_length + 50
)

tokenizer.decode(outputs[0])

実行

textを"日本について教えてください。"にした場合、出力は以下のようになりました。

<bos>日本について教えてください。\n\n日本はアジアの東部に位置する国です。アジアには多くの国があり、日本はアジアの東部に位置する国の一つです。\n\n日本は歴史的に重要な国です。古代の日本には、多くの伝説や神話があると考えられています。また、日本の歴史には、多くの戦国時代や時代劇があると考えられています。\n\n日本は経済的に重要な国です。日本は、世界で最も経済的な国の一つであり、経済的な成長は非常に高いです。\n\n日本は文化的に重要な国です

Tiny llamaとの比較

軽量LLMのTiny llamaと比較しました。Tiny llamaは日本語での入出力用に学習されていないため以降は英語の入力で比較を行っています。

Gemma 2B it Tiny llama
メモリ使用量 10.14GB 4.36GB
Generate時間 30.97秒 6.46秒

各モデルの出力は以下のようになりました。

Gemma 2B it

<bos>My name is Ken. I am 25 years old and I live in London, England. I am a software engineer and I have been working in the industry for 5 years.\n\nI am a passionate and dedicated individual with a strong work ethic and a thirst for knowledge. I am always looking for new challenges and opportunities to learn and grow.\n\nI am a good communicator and I am always willing to help others. I am also a team player and I am always willing to contribute to the success of the team.\n\nI am a big fan of technology and I am'

Tiny llama

<s> My name is Ken. I am 25 years old and I live in the United States. I am a student at the University of California, Berkeley. I am a member of the American Association of University Professors. I am a member of the American Association of University Professors. I am a member of the American Association of University Professors. I am a member of the American Association of University Professors. I am a member of the American Association of University Professors. I am a member of the American Association of University Professors.'

おわりに

GPUなしで動く高速、軽量モデルへの関心が高まれば、ローカルで動くアプリが生まれるのではないかと思いました。また、cppモデルによる高速化やMacGPUを使った高速化を試していきたいと思います。