Mambaを動かして速度をtransformerと比較するメモ

Kan Hatakeyama

はじめに

話題沸騰中のMambaを動かしてみました｡速度を同等サイズのtransformerと比較してみます｡

関連ページ

Colabで動かす

動かした結果

基本的には､npakaさんのコードを貼り付けただけです｡
無料のGPU(T4)を使いました｡

実行コードの部分はこちら｡
(ゆくゆく､コードの中身もチェックしたいです｡)


!python -i -m benchmarks.benchmark_generation_mamba_simple \
  --model-name "kotoba-tech/kotomamba-2.8B" \
  --tokenizer-path "kotoba-tech/kotomamba-2.8B" \
  --tokenizer-model "/root/.cache/huggingface/hub/models--kotoba-tech--kotomamba-2.8B-v1.0/snapshots/399c741bb23584bf4eb93805991c1cdb0dc0368a/tokenizer.model" \
  --tokenizer-type SentencePieceTokenizer \
  --use-sentencepiece \
  --prompt "Q: Who is Doraemon? Gime me detailed answers. A:" \
  --topp 0.9 \
  --temperature 0.7 \
  --repetition-penalty 1.2

実行結果
ドラえもんは誰ですか?と英語で聞いてみました｡ (質問のスペルミスってます)

['<s> Q: Who is Doraemon? Gime me detailed answers. A: He\'s a little boy who lives in the city of Paradiso with his parents and older brother, Kotaro (whom he calls "Ko"). His father works at an electronics company while their mother runs her own business as well; they have two children named Chibi-Robo ("Chi" for short) & Dorypido which means both boys are called by this name).\nQ: What does it mean to be born on Christmas Day or New Year´s Eve ? How do you celebrate them?? 10 questions about Japan!'] Prompt length: 16, generation length: 114 kotoba-tech/kotomamba-2.8B prompt processing + decoding time: 3061ms >>>

日本語訳

Q: ドラえもんって誰？詳しく教えてください。A：パラディソの街で、両親と兄のコタロー（コと呼ぶ）と一緒に暮らしている少年。ちびロボ（略して "ち"）とドリピド（略して "ドリ"）という2人の子供がいる。また、どのようにお祝いしますか？日本に関する10の質問！

パラディソの街で暮らしている少年だそうです｡誰なんですかね｡

因みに､日本語で聞いた結果はこちら｡

<s> Q: ドラえもんは誰ですか?詳しく教えて A:のび太です。\nQ.102：あなたが好きな漫画家を3人挙げなさい!A:赤塚不二夫、藤子・F・不二雄と石ノ森章太郎だよ!!

回答は間違ってますが､少しはマシになりました｡
(日英の知識転移はうまくいっていないみたいですね｡)

推論時間

3回の平均を取りました｡

time: 4.695, speed: 37.273695420660275 tokens/sec
time: 3.301, speed: 34.53498939715238 tokens/sec
time: 3.061, speed: 36.91604050963738 tokens/sec

T4 GPUでのMambaの生成速度は､30-40 token/secでした｡

transformerでのコントロール実験

phi-2(2.7b)を試す

同等サイズのtransformerモデルである､phi-2 (2.7b)の生成速度を調べました｡

import time
q="Q: Who is Doraemon? Gime me detailed answers. A:"

for i in range(3):
  start_time=time.time()
  res=pipe(q, max_new_tokens=512)[0]["generated_text"][len(q):]
  end_time=time.time()
  print(res)
  n_tokens=len(tokenizer.tokenize(res))
  consumed_time=end_time-start_time
  print(f"time consumed: {consumed_time}")
  print(n_tokens/(consumed_time)," tokens/sec")

回答は以下の通り｡

Doraemon is a robotic cat from the future. He has many gadgets and can help his friends with various problems.
(ドラえもんは未来からやってきたロボット猫。たくさんのガジェットを持ち、様々な問題で友達を助けることができる。)

正解ですね｡

生成速度

time consumed: 1.8966805934906006
13.180922547423055  tokens/sec

time consumed: 1.558398723602295
16.042107595039347  tokens/sec

time consumed: 1.5314698219299316
16.32418715799142  tokens/sec

今回のtransformerの場合は､平均で15 tokens/sec程度のようでした｡

まとめ

tokenizerの違いなど､諸々の細かな条件を話し出すと､キリがないですが､
今回のT4 GPUでの比較では､mamba (>30 tokens/sec)がtransformer(~ 15 tokens/sec)よりも2倍以上､早いという結果になりました｡

言語モデルの推論スピードは早いに越したことないので､mambaは有効な選択肢になりそうです｡


補足: Scaling law的にも､mambaの学習効率は悪く無さそうです｡欠点はどこにあるんでしょうか..｡



この記事が気に入ったらサポートをしてみませんか？