Cloudflare社のWorkers AIでも使用されているM2M100という多言語翻訳モデルを使ってみた

こんちには。

データアナリティクス事業本部 インテグレーション部 機械学習チームの中村です。

今回はMeta社が公開しているM2M100という多言語翻訳モデルについての使い方を見ていきます。

M2M100は多言語エンコーダ・デコーダ（seq-to-seq）モデルであり、100言語、9,900方向の翻訳処理が可能なモデルとなっています。Cloudflare社が提供するWorkers AIのTranslation処理も、このM2M100の1.2Bモデルがコアとなっているため、そこでM2M100が気になり使ってみることにしました。

M2M100には418M、1.2B、12Bと3種類のサイズがありますが、今回は1.2Bサイズを動かしてみます。

また、ただ動かすだけに加えて以下に記載されている最適化も試してみました。

実行環境

Google Colaboratoryで行い、ランタイムタイプは以下のようにしました。

• アクセラレータ : T4 GPU (GPU RAM 15.0GB)
• ハイメモリ : オフ (RAM 12.7GB)

実行環境のソフトウェアバージョンは以下となります。

• Python : 3.10.12

準備

以下のモジュールをインストールします。

!pip install transformers
!pip install sentencepiece
!pip install bitsandbytes

transformersとsentencepieceは、M2M100を動作させるために必要なモジュールです。

bitsandbytesは、省メモリ化、最適化するために必要なライブラリとなり、将来的には12Bモデルの方も試すことを念頭に置いて今回こちらも試してみました。

モジュールのバージョンは下記のようになりました。

!pip freeze | grep \
    -e "torch==" -e "torch" -e "transformers" \
    -e "sentencepiece" -e "bitsandbytes"

# bitsandbytes==0.41.1
# sentencepiece==0.1.99
# torch @ https://download.pytorch.org/whl/cu118/torch-2.0.1%2Bcu118-cp310-cp310-linux_x86_64.whl#sha256=a7a49d459bf4862f64f7bc1a68beccf8881c2fa9f3e0569608e16ba6f85ebf7b
# torchaudio @ https://download.pytorch.org/whl/cu118/torchaudio-2.0.2%2Bcu118-cp310-cp310-linux_x86_64.whl#sha256=26692645ea061a005c57ec581a2d0425210ac6ba9f923edf11cc9b0ef3a111e9
# torchdata==0.6.1
# torchsummary==1.5.1
# torchtext==0.15.2
# torchvision @ https://download.pytorch.org/whl/cu118/torchvision-0.15.2%2Bcu118-cp310-cp310-linux_x86_64.whl#sha256=19ca4ab5d6179bbe53cff79df1a855ee6533c2861ddc7389f68349d8b9f8302a
# transformers==4.33.3

翻訳する入力データ

Amazon Bedrockが話題となっていますので、その公式アナウンスの文章を入力データとして使ってみます。

• https://aws.amazon.com/jp/about-aws/whats-new/2023/09/amazon-bedrock-generally-available/

text = """
Amazon Bedrock, the easiest way to build and scale generative AI applications with foundation models (FMs), is now generally available.
Amazon Bedrock is a fully managed service that offers a choice of high-performing FMs from leading AI companies including AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon, along with a broad set of capabilities that you need to build generative AI applications, simplifying development while maintaining privacy and security.
Amazon Bedrock is the first to offer Llama 2, Meta’s large language models (LLMs), in fine-tuned 13B and 70B parameter versions as a fully managed API.
Llama models are ideal for dialogue use cases. To learn more, see Llama 2 on Amazon Bedrock.
To help you accelerate deploying generative AI into production, provisioned throughput is available in Amazon Bedrock, which provides you the flexibility and control to reserve throughput (Input/Output tokens per minute) and maintain a consistent user experience even during peak traffic times.
For customers building in highly regulated industries, Amazon Bedrock has achieved HIPAA eligibility and GDPR compliance.
Additionally, Amazon Bedrock is integrated with Amazon CloudWatch, to help you track usage metrics and build customized dashboards for audit purposes, and with AWS CloudTrail, to monitor and troubleshoot API activity as you integrate other systems into your generative AI applications.
Amazon Bedrock is available in the US East (N. Virginia) and US West (Oregon) AWS Regions.
To learn more about building with generative AI, visit Amazon Bedrock, Amazon Bedrock pricing, and the announcement blog post.
"""

これをセンテンス単位に分割して、翻訳処理を実施します。

sentences = [t for t in text.split("\n") if len(t)!=0]

推論用の関数

また翻訳モデル推論用の関数として以下を作成しておきます。

def translate(sentences, tokenizer, model, src_lang="en", target_lang="ja"):
    tokenizer.src_lang = src_lang
    encoded = tokenizer(sentences, return_tensors="pt", padding=True).to("cuda")
    generated_tokens = model.generate(**encoded,
        forced_bos_token_id=tokenizer.get_lang_id(target_lang), max_new_tokens=100)
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)

sentencesがセンテンス単位に分割されたセンテンスのリスト、tokenizerとmodelがモデル本体です。

（ちなみに１センテンス単位で分割して推論するのが最適化どうかは検討の余地があります）

src_langに翻訳元の言語を指定し、tokenizer.src_langにこれを設定しています。

tokenizerの処理でテキストを数値にエンコードする際に、padding=Trueを与えている点はポイントとなります。 今回は複数センテンスをバッチ処理するため、エンコード結果の長さが統一されるようpadding=Trueの設定が必要となります。

また基本的に処理をGPU実行するために、エンコード結果は.to("cuda")でGPU側に転送しています。

その他、各文に対して十分な長さで翻訳結果が得られるように、max_new_tokens=100を指定しています。

ダウンロード

モデルをはじめにダウンロードしておきます。

%%time
from huggingface_hub import snapshot_download
snapshot_download("facebook/m2m100_1.2B")

# CPU times: user 16 s, sys: 14.9 s, total: 30.9 s
# Wall time: 4min 17s

ダウンロードされたデータは約10GB程度となりました。

!du -BM /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B

# 1M      /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/refs
# 9463M   /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/blobs
# 1M      /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/snapshots/c24bcaa3d5101a2535d96f0029af4bebf124cfe6
# 1M      /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/snapshots
# 1M      /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/.no_exist/c24bcaa3d5101a2535d96f0029af4bebf124cfe6
# 1M      /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B/.no_exist
# 9463M   /root/.cache/huggingface/hub/models--facebook--m2m100_1.2B

コード

デフォルト (最適化なし)

まずは1.2Bモデルを動かしてみます。モデルのロード処理は以下です。

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B")\
    .to("cuda")

1.2Bサイズの場合は、このままで特に工夫なく読み込むことができます。

以下で翻訳処理を動かすことができます。

translated = translate(sentences, tokenizer, model)

以降、翻訳処理のコードは共通となるため割愛します。

また翻訳結果と処理時間の比較については後でまとめて記載します。

16bit量子化 (float16)

次は浮動小数点を16bitにしてみます。この場合のモデルのロード処理は以下です。

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
import torch

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B",
    torch_dtype=torch.float16)\
    .to("cuda")

16bit量子化 (bfloat16)

16bitの浮動小数点にはfloat16とは別にbfloat16というdata typeもあります。

こちらは、float16のビットサイズと同じにしつつ、float32相当のダイナミックレンジ（取りうる値の範囲）を持つように指数部と仮数部が調整されたものというイメージです。

以下の図が分かりやすいかと思います。

（出典：https://cloud.google.com/tpu/docs/bfloat16?hl=ja）

この場合のモデルのロード処理は以下です。

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
import torch

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B",
    torch_dtype=torch.bfloat16)\
    .to("cuda")

bitsandbytes (LLM.int8)

次にbitsandbytesによる量子化も試してみます。まずはLLM.int8()と呼ばれている手法で、Transformerにおけるfeed-forward層とattention projection層の行列演算を、整数8bit型で実行することにより、メモリ使用量を抑える手法となっています。

この場合のモデルのロード処理は以下です。

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
import torch

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100-12B-last-ckpt")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100-12B-last-ckpt",
    torch_dtype=torch.bfloat16, device_map="auto", load_in_8bit=True)

またBitsAndBytesConfigクラスを使えば、より細かいレベルで設定が可能です。（今回実行は割愛）

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
from transformers import BitsAndBytesConfig
import torch

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6,
    llm_int8_skip_modules=[],
    llm_int8_enable_fp32_cpu_offload=False,
    llm_int8_has_fp16_weight=False
)

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100-12B-last-ckpt")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100-12B-last-ckpt",
    torch_dtype=torch.bfloat16, quantization_config=quantization_config)

各パラメータの意味は以下に記載されています。

• Quantize 🤗 Transformers models

詳細は以下の論文やHugging Faceのブログを参照ください。

bitsandbytes (FP4)

bitsandbytesのもう一つのオプションとしてFP4が準備されており、データを4bitの浮動小数点で保存し演算時のみに16bitで計算することで、メモリ使用量を削減します。

この場合のモデルのロード処理は以下です。

%%time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
import torch

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100-12B-last-ckpt")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100-12B-last-ckpt",
    torch_dtype=torch.bfloat16, device_map="auto", load_in_4bit=True)

BitsAndBytesConfigを使えば、より細かいレベルで設定が可能です。（今回実行は割愛）

from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
from transformers import BitsAndBytesConfig
import torch

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4", # 推論時はnf4の恩恵はほぼないと記載されており学習時がメインのようです
    bnb_4bit_use_double_quant=True
)

tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100-12B-last-ckpt")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100-12B-last-ckpt",
    torch_dtype=torch.bfloat16, quantization_config=quantization_config)

各パラメータの意味は以下に記載されています。

• Quantize 🤗 Transformers models

この機能は比較的最近、transformersにマージされた機能のようです。

• 4-bit QLoRA via bitsandbytes (4-bit base model + LoRA) by TimDettmers · Pull Request #23479 · huggingface/transformers · GitHub

ベースはQLoRAの論文で使われる技術を元にされているようで、fine-tuning時に主に関連のあるオプションもあります。

詳細は以下のブログを参照ください。

結果

翻訳結果の比較

以下に各センテンスの比較結果を記載します。参考までにDeepLの結果も併記しています。

（なお、float16とbfloat16は、defaultと同じであったため今回記載を省略しました。）

定量評価はできませんが、どれも意味は捉えられる文章になっているかなと思います。

またdefaultとLLM.int8は同じ結果となることが多い一方、defaultとFP4はそれよりも異なる結果が多い印象です。

処理時間の比較

各モードでリソースの処理時間や使用量を確認しました。

16bit量子化は、推論時間的にもCPUメモリ、GPUメモリ使用量的にもメリットがありそうです。

LLM.int8は今回の場合は推論時間が遅く、CPUメモリ使用量が増加するため、あまりメリットがないような結果となりました。

FP4は、よりGPUメモリが小さくなりますが、CPUメモリ資料量は16bitより増え、推論時間も増加するため注意が必要です。

まとめ

いかがでしたでしょうか。M2M100モデルを題材に、その翻訳結果と各最適化オプションの推論時の影響を確認できました。

M2M100の1.2Bモデルでもある程度意味の合っている翻訳結果となることが分かりました。英語にしか対応してないモデルの前処理としても検討の余地があるかもしれません。

最適化オプションの結果は、他の入力データやモデルでも必ず同じ結果となるわけではなく、12Bサイズのモデルだとまた違った結果となるかもしれませんので、その点はご留意ください。

本記事が皆様のご参考になれば幸いです。