Quantization
littlemexSageMaker での量子化プロセスの流れ
量子化プロセスの詳細説明
量子化技術の比較
量子化のメモリ削減効果
技術的な詳細説明
SageMaker の Hugging Face LLM コンテナ内部では、以下のプロセスが実行されています：
-
モデルのダウンロード:
• Hugging Face Hub から指定されたモデル（例: tokyotech-llm/Llama-3.3-Swallow-70B-v0.4）をダウンロード
• モデルのアーキテクチャとトークナイザーを読み込み -
量子化の適用:
• HF_MODEL_QUANTIZE 環境変数に基づいて量子化方式を決定
• 選択された量子化ライブラリを初期化
• モデルの重みを指定された精度（例: 4ビット）に変換 -
bitsandbytes-nf4 の場合:
• モデルの重みを Normal Float 4 (NF4) 形式に変換
• 一部の重要なレイヤー（例: 出力層）は FP16/BF16 精度を維持
• 量子化された重みを GPU メモリにロード -
gptq の場合:
• 事前に計算された量子化パラメータを使用（または実行時に計算）
• モデルの重みを GPTQ 形式に変換
• カスタムカーネルを使用して高速な推論を実現 -
awq の場合:
• アクティベーション値を考慮した量子化を適用
• モデルの重みを AWQ 形式に変換
• 精度を最大限に保ちながらメモリ使用量を削減 -
Text Generation Inference サーバーの起動:
• 量子化されたモデルを Text Generation Inference (TGI) サーバーに渡す
• バッチ処理、キャッシュ、その他の最適化を適用
• REST API エンドポイントを公開 -
SageMaker エンドポイントの準備:
• ヘルスチェックを実行して TGI サーバーが正常に動作していることを確認
• エンドポイントを「InService」状態に設定
この一連のプロセスにより、元の大規模モデルが量子化され、限られたメモリリソースでも効率的に実行できるようになります。
littlemex
Transformer Engine での FP8 の使用
H100 GPU は新しいデータ型である FP8（8ビット浮動小数点）のサポートを導入し、行列乗算と畳み込みの高いスループットを実現しています。この例では、FP8 データ型を紹介し、Transformer Engine での使用方法を示します。
FP8 の概要
構造
H100 GPU がサポートする FP8 データ型は、ニューラルネットワークのトレーニングの異なる部分で有用な 2 つの異なるデータ型です：
-
E4M3 - 1 つの符号ビット、4 つの指数ビット、3 つの仮数ビットで構成されています。+/-448 および
nanまでの値を格納できます。 -
E5M2 - 1 つの符号ビット、5 つの指数ビット、2 つの仮数ビットで構成されています。+/-57344、+/-
inf、およびnanまでの値を格納できます。動的範囲が増加する代わりに、格納される値の精度が低下します。
図 1: 浮動小数点データ型の構造。表示されているすべての値（FP16、BF16、FP8 E4M3、FP8 E5M2）は、値 0.3952 の最も近い表現です。
ニューラルネットワークのトレーニング中、これらの型の両方が利用される場合があります。一般的に、順伝播の活性化と重みはより高い精度を必要とするため、順伝播では E4M3 データ型が最適です。しかし、逆伝播では、ネットワークを流れる勾配は通常、精度の損失の影響を受けにくいですが、より高い動的範囲を必要とします。したがって、E5M2 データ形式を使用して格納するのが最適です。H100 TensorCore は、これらの型の任意の組み合わせを入力としてサポートしており、各テンソルを好みの精度で格納することができます。
混合精度トレーニング - 簡単な紹介
FP8 がディープラーニングモデルのトレーニングにどのように使用できるかを理解するために、まず他のデータ型、特に FP16 での混合精度の仕組みを思い出すと役立ちます。
FP16 トレーニングの混合精度レシピには 2 つの要素があります：FP16 で実行する操作の選択と動的損失スケーリングです。
-
FP16 精度で実行する操作の選択は、操作の入力に対する出力の数値的な挙動と、期待されるパフォーマンス向上の分析が必要です。これにより、行列乗算、畳み込み、正規化レイヤーなどの操作を安全とマークし、
normやexpなどの操作は高精度を必要とするものとして残すことができます。 -
動的損失スケーリングは、トレーニング中の勾配のオーバーフローとアンダーフローの両方を回避することを可能にします。FP16 の動的範囲は勾配値の分布を格納するのに十分ですが、この分布が FP16 が処理するには高すぎるか低すぎる値を中心に配置されている場合に発生する可能性があります。損失をスケーリングすることで、これらの分布を（2 のべき乗のみを使用して数値に影響を与えずに）FP16 で表現可能な範囲にシフトします。
図 2: 損失をスケーリングすることで、勾配分布を FP16 データ型の表現可能な範囲にシフトできます。
FP8 での混合精度トレーニング
FP8 型が提供する動的範囲は、特定の活性化または勾配を格納するには十分ですが、同時にすべてを格納するには十分ではありません。これにより、FP16 で機能した単一の損失スケーリング係数戦略が FP8 トレーニングでは実行不可能となり、代わりに各 FP8 テンソルに対して異なるスケーリング係数を使用する必要があります。
特定の FP8 テンソルに適切なスケーリング係数を選択するための複数の戦略があります：
-
ジャストインタイムスケーリング。この戦略は、生成されるテンソルの絶対値の最大値（amax）に基づいてスケーリング係数を選択します。実際には、データを複数回通過する必要があるため実行不可能です - オペレータは出力を生成し、より高い精度で書き出し、次に出力の絶対値の最大値を見つけて、すべての値に適用して最終的な FP8 出力を得ます。これにより多くのオーバーヘッドが発生し、FP8 使用によるゲインが大幅に減少します。
-
遅延スケーリング。この戦略は、以前の反復で見られた絶対値の最大値に基づいてスケーリング係数を選択します。これにより FP8 計算の完全なパフォーマンスが可能になりますが、最大値の履歴を FP8 オペレータの追加パラメータとして保存する必要があります。
図 3: 遅延スケーリング戦略。FP8 オペレータは、以前の反復で見られた amax（絶対値の最大値）の履歴を使用して得られたスケーリング係数を使用し、FP8 出力と現在の amax の両方を生成し、amax は履歴に保存されます。
図 3 に見られるように、遅延スケーリング戦略は amax の履歴を保存するだけでなく、その履歴を次の反復で使用されるスケーリング係数に変換するためのレシピを選択する必要もあります。
MXFP8 とブロックスケーリング
NVIDIA Blackwell アーキテクチャは、FP8 フォーマットの新しいバリアントである MXFP8 のサポートを導入しました。
MXFP8 vs FP8
「通常の」FP8 と MXFP8 の主な違いは、スケーリングの粒度にあります。FP8 では、各テンソルには単一の FP32 スケーリング係数があるため、テンソル内のすべての値が FP8 データ型の動的範囲内に「収まる」必要があります。これにより、ネットワーク内の一部のテンソル（勾配など）を表現するために、より精度の低い E5M2 フォーマットを使用する必要があります。
MXFP8 は、32 の連続した値の各ブロックに異なるスケーリング係数を割り当てることでこの問題に対処します。これにより、すべての値を E4M3 データ型で表現することができます。
図 4: MXFP8 は単一のテンソルに対して複数のスケーリング係数を使用します。図ではわかりやすくするためにブロックあたり 4 つの値しか示していませんが、実際の MXFP8 はブロックあたり 32 の値を持ちます。
図 5: 複数のスケーリング係数により、テンソルの動的範囲要件が減少し、E4M3 フォーマットを使用できるようになります。これは、0 に飽和する要素が少なくなるためです。
2 つ目の違いは、スケーリング係数を格納するために使用されるデータ型です。FP8 は FP32（E8M23）を使用しますが、MXFP8 は 2 のべき乗の 8 ビット表現（E8M0）を使用します。
図 6: MXFP8 でスケーリング係数を格納するために使用される E8M0 データ型の構造。
転置の処理
線形層の順伝播と逆伝播には、異なる縮約次元を持つ複数の行列乗算が含まれます。Blackwell Tensor Core は、MXFP8 データが縮約次元に対して「連続」であることを要求するため、MXFP8 トレーニングでは異なるポイントで非転置および転置された MXFP8 テンソルを使用します。ただし、FP8 データの転置は数値的に自明ですが、MXFP8 データの転置には再量子化が必要です。
この二重量子化に関連する精度の損失を避けるために、Transformer Engine は元の高精度入力からテンソルの通常のコピーと転置されたコピーの両方を作成します。
図 7: MXFP8 での線形層。順伝播と逆伝播の両方を計算するには、両方の方向で量子化されたテンソルが必要です。
Transformer Engine での FP8 の使用
Transformer Engine ライブラリは、FP8 遅延スケーリングと MXFP8 戦略を使用した FP8 データ型でのトレーニングを容易に使用するためのツールを提供します。
FP8 レシピ
transformer_engine.common.recipe モジュールの DelayedScaling レシピは、FP8 遅延スケーリングでのトレーニングに必要なすべてのオプションを格納します：スケーリング係数計算に使用する amax 履歴の長さ、FP8 データフォーマットなど。同様に、同じモジュールの MXFP8BlockScaling を使用して MXFP8 トレーニングを有効にすることができます。
from transformer_engine.common.recipe import Format, DelayedScaling, MXFP8BlockScaling
fp8_format = Format.HYBRID # 順伝播中は E4M3、逆伝播中は E5M2
fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")
mxfp8_format = Format.E4M3 # E4M3 をすべての場所で使用
mxfp8_recipe = MXFP8BlockScaling(fp8_format=mxfp8_format)
このレシピは、FP8 トレーニングを設定するために使用されます。
FP8 オートキャスティング
すべての操作が FP8 を使用して安全に実行できるわけではありません。Transformer Engine ライブラリによって提供されるすべてのモジュールは、精度を維持しながら FP8 データ型から最大のパフォーマンス利益を提供するように設計されています。FP8 操作を有効にするには、TE モジュールを fp8_autocast コンテキストマネージャ内でラップする必要があります。
import transformer_engine.pytorch as te
import torch
torch.manual_seed(12345)
my_linear = te.Linear(768, 768, bias=True)
inp = torch.rand((1024, 768)).cuda()
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
out_fp8 = my_linear(inp)
fp8_autocast コンテキストマネージャは FP8 の処理の複雑さを隠します：
- すべての FP8 安全操作は入力を FP8 にキャストします
- amax 履歴が更新されます
- 新しいスケーリング係数が計算され、次の反復の準備ができています
注意
Transformer Engine の Linear レイヤーでの FP8 のサポートは現在、両方の次元が 16 で割り切れる形状のテンソルに限定されています。完全な Transformer ネットワークへの入力に関しては、これは通常、シーケンス長を 16 の倍数にパディングする必要があることを意味します。
逆伝播の処理
モデルが fp8_autocast 領域内で実行される場合、特にマルチ GPU トレーニングでは、スケーリング係数と amax 履歴を同期するために一部の通信が必要です。多くのオーバーヘッドを導入せずにその通信を実行するために、fp8_autocast コンテキストマネージャは通信を実行する前にテンソルを集約します。
この集約のため、逆伝播呼び出しは fp8_autocast コンテキストマネージャの外部で行う必要があります。これは計算精度に影響を与えません - 逆伝播の精度は順伝播の精度によって決まります。
loss_fp8 = out_fp8.mean()
loss_fp8.backward() # この逆伝播は FP8 を使用します。out_fp8 が fp8_autocast 内で計算されたため
out_fp32 = my_linear(inp)
loss_fp32 = out_fp32.mean()
loss_fp32.backward() # この逆伝播は FP8 を使用しません。out_fp32 が fp8_autocast 外で計算されたため
精度
FP32 と FP8 の実行結果を比較すると、比較的近いですが異なることがわかります：
out_fp8
tensor([[ 0.2276, 0.2627, 0.3001, ..., 0.0346, 0.2211, 0.1188],
[-0.0963, -0.3725, 0.1717, ..., 0.0901, 0.0522, -0.3472],
[ 0.4526, 0.3482, 0.5976, ..., -0.0687, -0.0382, 0.1566],
...,
[ 0.1698, 0.6061, 0.0385, ..., -0.2875, -0.1152, -0.0260],
[ 0.0679, 0.2946, 0.2751, ..., -0.2284, 0.0517, -0.1441],
[ 0.1865, 0.2353, 0.9172, ..., 0.1085, 0.1135, 0.1438]],
device='cuda:0', grad_fn=<_LinearBackward>)
out_fp32
tensor([[ 0.2373, 0.2674, 0.2980, ..., 0.0233, 0.2498, 0.1131],
[-0.0767, -0.3778, 0.1862, ..., 0.0858, 0.0676, -0.3369],
[ 0.4615, 0.3593, 0.5813, ..., -0.0779, -0.0349, 0.1422],
...,
[ 0.1914, 0.6038, 0.0382, ..., -0.2847, -0.0991, -0.0423],
[ 0.0864, 0.2895, 0.2719, ..., -0.2388, 0.0772, -0.1541],
[ 0.2019, 0.2275, 0.9027, ..., 0.1022, 0.1300, 0.1444]],
device='cuda:0', grad_fn=<_LinearBackward>)
これは、FP8 の場合、計算前に入力と重みの両方が FP8 にキャストされるために発生します。FP8 で表現可能な入力を使用すると、この違いを確認できます（quickstart_utils.py で定義された関数を使用）：
from quickstart_utils import cast_to_representable
inp_representable = cast_to_representable(inp)
my_linear.weight.data = cast_to_representable(my_linear.weight.data)
out_fp32_representable = my_linear(inp_representable)
print(out_fp32_representable)
tensor([[ 0.2276, 0.2629, 0.3000, ..., 0.0346, 0.2211, 0.1188],
[-0.0963, -0.3724, 0.1717, ..., 0.0901, 0.0522, -0.3470],
[ 0.4526, 0.3479, 0.5976, ..., -0.0686, -0.0382, 0.1566],
...,
[ 0.1698, 0.6062, 0.0385, ..., -0.2876, -0.1152, -0.0260],
[ 0.0679, 0.2947, 0.2750, ..., -0.2284, 0.0516, -0.1441],
[ 0.1865, 0.2353, 0.9170, ..., 0.1085, 0.1135, 0.1438]],
device='cuda:0', grad_fn=<_LinearBackward>)
今回の差はとても小さいです：
out_fp8 - out_fp32_representable
tensor([[ 4.9591e-05, -1.9073e-04, 9.5367e-05, ..., -3.8147e-06,
4.1962e-05, 2.2888e-05],
[ 2.2888e-05, -3.4332e-05, 2.2888e-05, ..., 2.6703e-05,
5.3406e-05, -1.4114e-04],
[-3.8147e-05, 2.6703e-04, -3.8147e-06, ..., -5.7220e-05,
4.1962e-05, -1.9073e-05],
...,
[ 1.1444e-05, -7.2479e-05, -3.8147e-06, ..., 5.3406e-05,
-1.5259e-05, 2.2888e-05],
[ 4.9591e-05, -9.5367e-05, 6.8665e-05, ..., -1.5259e-05,
7.6294e-05, 4.5776e-05],
[-1.5259e-05, -7.6294e-06, 1.8692e-04, ..., -3.0518e-05,
-4.5776e-05, 7.6294e-06]], device='cuda:0', grad_fn=<SubBackward0>)
FP8 実行から生じる結果の違いはトレーニングプロセス中には問題になりませんが、例えばモデルのデバッグ中などに理解しておくと良いでしょう。
littlemexでは、小さな具体例で説明します：
元のFP32データ（重み行列）があるとして：
W = [2.34 1.89]
[0.75 3.42]
- 最初から2つのMXFP8バージョンを作成：
通常版（一回の量子化）：
W_normal = [2.375 1.875] // MXFP8に量子化
[0.750 3.375]
転置版（一回の量子化）：
W_transposed = [2.375 0.750] // 転置してから量子化
[1.875 3.375]
問題のある方法（二重量子化）：
W_normal = [2.375 1.875] // まずMXFP8に量子化
[0.750 3.375]
↓ 転置して再量子化
W_double_quantized = [2.375 0.750] // 精度がさらに落ちる
[1.875 3.375]
このように、最初から2つのバージョンを作ることで、二重量子化による追加の誤差を防いでいます。
実際の計算では、順伝播でW_normalを使い、逆伝播でW_transposedを使用することで、それぞれ一回の量子化誤差で済みます。
（注：実際のMXFP8の量子化値は、より複雑なスケーリングと量子化プロセスを使用します。この例は概念を示すための簡略化されたものです）
littlemexBlackwell
FP6, FP4 対応追加
Block FP16
共通指数部 + 仮数部でデータを表現する
(16 [1.5, 0.5...]) -> (24, 8)