TensorFlowとXLAを使用した高速なテキスト生成

TensorFlowとXLAを使った高速なテキスト生成

TL;DR : TensorFlowを使用した🤗transformersを使ったテキスト生成は、XLAでコンパイルできるようになりました。これにより、以前よりも最大100倍高速化され、PyTorchよりもさらに高速になりました – 以下のコラボをチェックしてください！

テキスト生成

大規模言語モデルの品質が向上するにつれて、そのモデルができることに対する私たちの期待も高まりました。特にOpenAIのGPT-2のリリース以来、テキスト生成機能を持つモデルが注目されています。そして、その理由は妥当です – これらのモデルは要約、翻訳、さらにはいくつかの言語タスクでのゼロショット学習能力のデモンストレーションを行うことができます。このブログ記事では、TensorFlowを使用してこのテクノロジーを最大限に活用する方法を紹介します。

🤗transformersライブラリはNLPモデルから始まりましたので、テキスト生成は私たちにとって非常に重要な要素です。アクセス可能で、簡単に制御可能で効率的であることをHugging Faceの民主化の取り組みの一環として保証することが目的です。テキスト生成のさまざまなタイプについては以前のブログ記事があります。ただし、以下にコア機能のクイックな概要があります – 私たちのgenerate関数に慣れている場合や、TensorFlowの特異性に直接ジャンプしたい場合は、スキップしても構いません。

まずは基本から始めましょう。テキスト生成は、do_sampleフラグによって確定的または確率的になります。デフォルトでは、Falseに設定されており、出力は確定的であるため、Greedy Decodingとも呼ばれます。それをTrueに設定すると、サンプリングとも呼ばれるため、出力は確率的になりますが、seed引数（stateless TensorFlowのランダム数生成と同じ形式）を介して再現可能な結果を得ることもできます。一般的なガイドラインとして、モデルから事実情報を得る場合は確定的な生成を、よりクリエイティブな出力を目指す場合は確率的な生成を望むでしょう。

• Nyström形式：ニュストローム法による線形時間とメモリでのセルフアテンションの近似
• プロキシマルポリシーオプティマイゼーション（PPO）
• Sentence Transformersモデルのトレーニングと微調整

# transformers >= 4.21.0が必要です
# サンプリングの出力は、使用するハードウェアによって異なる場合があります。
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
inputs = tokenizer(["TensorFlow is"], return_tensors="tf")

generated = model.generate(**inputs, do_sample=True, seed=(42, 0))
print("サンプリングの出力：", tokenizer.decode(generated[0]))
# > サンプリングの出力：TensorFlowは、データサイエンスでデータ構造と構造について学ぶための素晴らしい学習プラットフォームです。

ターゲットアプリケーションによっては、より長い出力が望ましい場合もあります。テキスト生成の出力の長さは、max_new_tokensを使用して制御できますが、より長い生成はより多くのリソースを必要とすることに注意してください。

generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), max_new_tokens=5
)
print("5つの新しいトークンに制限：", tokenizer.decode(generated[0]))
# > 5つの新しいトークンに制限：TensorFlowは、学習プラットフォームとして優れたものです。
generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), max_new_tokens=30
)
print("30つの新しいトークンに制限：", tokenizer.decode(generated[0]))
# > 30つの新しいトークンに制限：TensorFlowは、データサイエンスでデータ構造と構造について学ぶための素晴らしい学習プラットフォームです。

サンプリングにはいくつかの設定があり、ランダム性を制御するために調整できます。最も重要なのはtemperatureです。これは出力の全体的なエントロピーを設定します – 1.0以下の値は、より高い確率でトークンをサンプリングすることを優先します。一方、1.0を超える値は逆の効果をもたらします。それを0.0に設定すると、動作はGreedy Decodingになります。非常に大きな値は一様なサンプリングに近似します。

generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), temperature=0.7
)
print("温度0.7：", tokenizer.decode(generated[0]))
# > 温度0.7：TensorFlowは、これを行うための素晴らしい方法です........
generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), temperature=1.5
)
print("温度1.5：", tokenizer.decode(generated[0]))
# > 温度1.5：TensorFlowはCythonとBambooの両方で開発されています。
# Bamboo上で...

サンプリングとは対照的に、グリーディデコーディングは常に各生成イテレーションで最も確率の高いトークンを選択します。しかし、しばしば最適でない出力をもたらします。結果の品質を向上させるためには、num_beams引数を使用できます。この引数が1より大きい場合、ビームサーチがトリガされ、高確率のシーケンスを継続的に探索します。この探索には追加のリソースと計算時間がかかります。

generated = model.generate(**inputs, num_beams=2)
print("ビームサーチの出力:", tokenizer.decode(generated[0]))
# > ビームサーチの出力: TensorFlowはオープンソース、オープンソース、
# 分散ソースのアプリケーションフレームワークです

最後に、サンプリングまたはビームサーチを実行する際には、num_return_sequencesを使用して複数のシーケンスを返すことができます。サンプリングの場合、それは同じ入力プロンプトから複数回生成を実行することと同等です。一方、ビームサーチの場合、降順で最もスコアの高い生成されたビームを返します。

generated = model.generate(**inputs, num_beams=2, num_return_sequences=2)
print(
    "生成された仮説全体:",
    "\n".join(tokenizer.decode(out) for out in generated)
)
# > 生成された仮説全体: TensorFlowはオープンソース、オープンソース、
# 分散ソースのアプリケーションフレームワークです
# > TensorFlowはオープンソース、オープンソース、分散ソースのアプリケーション
# フレームワークを可能にするものです

テキスト生成の基本は、制御が容易であることがわかります。しかし、上記の例ではカバーされていない多くのオプションがあり、高度な使用例についてのドキュメントを読むことをお勧めします。残念ながら、TensorFlowでgenerateを実行すると、実行に時間がかかることに気付くかもしれません。ターゲットアプリケーションが低レイテンシーまたは大量の入力プロンプトを必要とする場合、TensorFlowでテキスト生成を実行することは高コストな作業のように見えます。 😬

しかし心配はいりません！このブログ記事の残りの部分では、わずか1行のコードで劇的な改善ができることを示すことを目指しています。すぐにアクションに移りたい場合は、コラボには試していただけるインタラクティブな例があります！

TensorFlowとXLA

XLA、またはAccelerated Linear Algebra、は元々TensorFlowモデルの高速化を目的として開発されたコンパイラです。今日では、JAXの背後にあるコンパイラでもあり、PyTorchとも使用することができます。”コンパイラ”という言葉は、一部の人にとっては困難に聞こえるかもしれませんが、XLAはTensorFlowと簡単に使用できます。TensorFlowライブラリにパッケージ化されており、グラフを作成する関数のjit_compile引数でトリガーすることができます。

TensorFlow 1に精通している方にとっては、TensorFlowグラフの概念は自然なものです。まず、グラフを作成するために宣言的な方法で操作を定義します。その後、グラフを通じて入力を送り、出力を観察することができます。高速で効率的ですが、デバッグは困難です。TensorFlow 2では、Eager Executionが導入され、モデルを命令形式でコーディングできるようになりました。TensorFlowチームは、その違いについて詳しく説明しています。

Hugging Faceは、自分たちのTensorFlowモデルをEager Executionを意識して作成しています。透明性は核となる価値であり、モデルの内部をいつでも検査できることは非常に有益です。ただし、それはモデルの一部の使用法がデフォルトではグラフモードのパフォーマンスの利点を得ることができないことを意味します（たとえば、model(args)を呼び出す場合など）。

幸いにも、TensorFlowチームは私たちのようなユーザーをカバーしてくれています 🥳！TensorFlowコードを含む関数をtf.functionでラップすると、ラップされた関数を呼び出すときにそれをグラフに変換しようとします。モデルをトレーニングしている場合、model.compile()（run_eagerly=Trueなし）を呼び出すと、それがラッピングを行い、model.fit()を呼び出すときにグラフモードの利点を得るために使用されます。TensorFlowコードを含む任意の関数でtf.functionを使用できるため、モデル推論を超える関数にも使用でき、最適化された単一のグラフを作成することができます。

TensorFlowグラフを作成する方法を知ったので、XLAでそれらをコンパイルすることは簡単です。上記の関数（tf.functionおよびtf.keras.Model.compile）にjit_compile=Trueという引数を追加するだけです。すべてがうまくいった場合（以下で詳しく説明します）、GPUまたはTPUを使用している場合、最初の呼び出しには時間がかかりますが、残りの呼び出しははるかに高速になります。以下は、モデルの推論とその出力のいくつかの後処理を行う関数の単純な例です：

# 注意：実行時間はハードウェアに深く依存します -- ここでは3090が使用されています。
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer(["TensorFlowは"], return_tensors="tf")

def most_likely_next_token(inputs):
    model_output = model(inputs)
    return tf.argmax(model_output.logits[:, -1, :], axis=-1)

print("TensorFlowのコードを使って通常の関数を呼び出しています...")
most_likely_next_token(inputs)
# > 実行時間 -- 48.8 ms

1行で上記の関数からXLAによって高速化された関数を作成することができます。

xla_most_likely_next_token = tf.function(most_likely_next_token, jit_compile=True)

print("XLA関数を呼び出しています...（初回は遅くなります）")
xla_most_likely_next_token(inputs)
# > 実行時間 -- 3951.0 ms
print("XLA関数を呼び出しています...（2回目は高速です）")
xla_most_likely_next_token(inputs)
# > 実行時間 -- 1.6 ms

XLAを使用したTensorFlowによるテキスト生成

最適化手法と同様に、XLAも例外ではありません。テキスト生成のユーザーの観点からは、心に留めておく必要がある技術的な側面が1つだけあります。詳細にはあまり立ち入らずに説明すると、この方法で使用されるXLAは、tf.functionのジャストインタイム（JIT）コンパイルを呼び出す際にポリモーフィズムに依存しています。

この方法で関数をコンパイルすると、XLAはすべてのテンソルの形状と型、および非テンソル関数入力のデータを追跡します。関数はバイナリにコンパイルされ、同じテンソル形状と型（どのテンソルデータでも）および同じ非テンソル引数で呼び出されるたびに、コンパイルされた関数を再利用できます。一方で、入力テンソルの形状や型が異なる場合、または異なる非テンソル引数を使用する場合は、新しいコストの高いコンパイルステップが行われます。簡単な例でまとめると次のようになります：

# 注意：実行時間はハードウェアに深く依存します -- ここでは3090が使用されています。
import tensorflow as tf

@tf.function(jit_compile=True)
def max_plus_constant(tensor, scalar):
    return tf.math.reduce_max(tensor) + scalar

# 遅い：XLAコンパイルがキックされるため、最初の呼び出しとなります
max_plus_constant(tf.constant([0, 0, 0]), 1)
# > 実行時間 -- 520.4 ms

# 速い：テンソルの形状、型、およびまったく同じ非テンソル引数での最初の呼び出しではありません
max_plus_constant(tf.constant([1000, 0, -10]), 1)
# > 実行時間 -- 0.6 ms

# 遅い：異なるテンソル型
max_plus_constant(tf.constant([0, 0, 0], dtype=tf.int64), 1)
# > 実行時間 -- 27.1 ms

# 遅い：異なるテンソル形状
max_plus_constant(tf.constant([0, 0, 0, 0]), 1)
# > 実行時間 -- 25.5 ms

# 遅い：異なる非テンソル引数
max_plus_constant(tf.constant([0, 0, 0]), 2)
# > 実行時間 -- 24.9 ms

実際のテキスト生成では、入力は特定の長さの倍数にパディングする必要があります（可能な形状の数が限られるため）、さらに異なるオプションを使用すると最初に遅くなることになります。XLAを使用して生成を単純に呼び出すとどうなるか見てみましょう。

# 注意：実行時間はハードウェアに深く依存します -- ここでは3090が使用されています。
import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# 新しい引数に注目してください、`padding_side="left"` -- TFAutoModelForCausalLMでインスタンス化できるデコーダーのみのモデルは、
# 入力プロンプトから生成を継続するため、左側にパディングする必要があります。
tokenizer = AutoTokenizer.from_pretrained(
    "gpt2", padding_side="left", pad_token="</s>"
)
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
input_1 = ["TensorFlowは"]
input_2 = ["TensorFlowはア"]

# 1行でXLA生成関数を作成します
xla_generate = tf.function(model.generate, jit_compile=True)

# パディングせずにXLA生成を呼び出します
tokenized_input_1 = tokenizer(input_1, return_tensors="tf")  # 長さ = 4
tokenized_input_2 = tokenizer(input_2, return_tensors="tf")  # 長さ = 5
print(f"`tokenized_input_1`の形状 = {tokenized_input_1.input_ids.shape}")
print(f"`tokenized_input_2`の形状 = {tokenized_input_2.input_ids.shape}")

print("tokenized_input_1でXLA生成を呼び出しています...")
print("(初回呼び出しのため遅くなります)")
start = time.time_ns()
xla_generate(**tokenized_input_1)
end = time.time_ns()
print(f"実行時間 -- {(end - start) / 1e6:.1f} ms\n")
# > 実行時間 -- 9565.1 ms

print("tokenized_input_2でXLA生成を呼び出しています...")
print("(異なる長さのため、トレースが再度トリガされます)")
start = time.time_ns()
xla_generate(**tokenized_input_2)
end = time.time_ns()
print(f"実行時間 -- {(end - start) / 1e6:.1f} ms\n")
# > 実行時間 -- 6815.0 ms

あら、それはとても遅いですね！さまざまな形状の組み合わせを管理するための解決策は、上記で言及されたように、パディングを使用することです。トークナイザークラスには、トレースを制限するために使用できるpad_to_multiple_of引数があります。

padding_kwargs = {"pad_to_multiple_of": 8, "padding": True}
tokenized_input_1_with_padding = tokenizer(
    input_1, return_tensors="tf", **padding_kwargs
)  # 長さ = 8
tokenized_input_2_with_padding = tokenizer(
    input_2, return_tensors="tf", **padding_kwargs
)  # 長さ = 8
print(
    "`tokenized_input_1_with_padding`の形状 = ",
    f"{tokenized_input_1_with_padding.input_ids.shape}"
)
print(
    "`tokenized_input_2_with_padding`の形状 = ",
    f"{tokenized_input_2_with_padding.input_ids.shape}"
)

print("tokenized_input_1_with_paddingを使用してXLA生成を呼び出しています...")
print("(遅いです、この長さで最初に実行する場合)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding)
end = time.time_ns()
print(f"実行時間 -- {(end - start) / 1e6:.1f} ms\n")
# > 実行時間 -- 6815.4 ms

print("tokenized_input_2_with_paddingを使用してXLA生成を呼び出しています...")
print("(速くなります！)")
start = time.time_ns()
xla_generate(**tokenized_input_2_with_padding)
end = time.time_ns()
print(f"実行時間 -- {(end - start) / 1e6:.1f} ms\n")
# > 実行時間 -- 19.3 ms

それはずっと良いですね。この方法で行われる連続した生成呼び出しは、以前のものよりも桁違いに高速です。ただし、新しい生成オプションを試すと、いつでもトレースがトリガーされることに注意してください。

print("同じ入力でXLA生成を呼び出していますが、新しいオプションを使用して...")
print("(再び遅いです)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding, num_beams=2)
end = time.time_ns()
print(f"実行時間 -- {(end - start) / 1e6:.1f} ms\n")
# > 実行時間 -- 9644.2 ms

開発者の観点からは、XLAを利用することはいくつかの追加的なニュアンスに注意することを意味します。モデルのトレーニングなど、データ構造のサイズが事前にわかっている場合、XLAは輝きます。一方、その次元が不可能であるか、特定の動的なスライスが使用される場合、XLAはコンパイルできません。現代のテキスト生成の実装は自己回帰的であり、テンソルを拡張し、途中でいくつかの操作を急停止するのが自然な動作です。つまり、デフォルトではXLAに対応していません。私たちは、テキスト生成のコード全体を書き直し、操作をベクトル化し、パディング付きの固定サイズの構造を使用するようにしました。また、NLPモデルも、パディングされた構造の存在下で位置埋め込みを正しく使用するように変更しました。結果として、TensorFlowのテキスト生成ユーザーには見えないはずですが、XLAのコンパイルが利用可能になります。

ベンチマークと結論

上記では、TensorFlowの関数をグラフに変換し、XLAのコンパイルでそれらを高速化できることを示しました。現在のテキスト生成の形式は、モデルの順方向パスといくつかの後処理を交互に繰り返し、イテレーションごとに1つのトークンを生成する自己回帰的な関数です。XLAのコンパイルにより、プロセス全体が最適化され、実行が高速化されます。しかし、どれくらい速くなるのでしょうか？以下のGradioデモでは、Hugging Faceのテキスト生成をTensorFlowとPyTorchの2つの主要なMLフレームワークで複数のGPUモデルで比較したベンチマークが含まれています。

結果を探索すると、2つの結論がすぐに明らかになります：

1. このブログ投稿がここまで構築されてきたように、XLAを使用するとTensorFlowのテキスト生成ははるかに高速です。コンパイルされたグラフを使用した場合、速度が100倍以上向上する場合もあります 🚀
2. XLAを使用したTensorFlowのテキスト生成は、ほとんどの場合で最も高速なオプションです。いくつかの場合では、9倍以上も速くなり、PyTorchが真剣なNLPタスクに最適なフレームワークであるという神話を打ち砕いています 💪

Colabを試してみて、XLAでスーパーチャージされたテキスト生成のパワーをお楽しみください！

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!