Transformers、量子化重みの読み込み処理を再設計　展開を読み込み時に前倒し

Hugging Face TransformersのコントリビューターであるArthurZuckerは11月2日、compressed-tensorsチェックポイントの読み込み方法を変更するPR #50238を提出しました。重み変換マッピングを通じて展開することで、デフォルトでは読み込み後に密な重みを取得できるようにし、初回のフォワード演算まで展開を待たない仕組みです。これは現在レビュー待ちの提案であり、正式リリースでの動作と見なすことはできません。提案の説明

この変更は、モデルファイルの形式と実行時の精度の違いに関わります。公式ドキュメントによると、compressed-tensorsはFP8やINT4などの圧縮重みを保存できます。現在のデフォルトのフローでは圧縮データを保持し、初回のフォワード演算時に展開します。ユーザーはdequantize=Trueを設定して、読み込み時にBF16などのモデルデータ型へ復元することもできます。FP8を維持して最適化カーネルを使うには、use_optimized_inference=Trueを明示的に有効にする必要があります。読み込みモードのドキュメント

提案が具体的に対処するのは、読み込み時にチェックポイント内の重み名を再マッピングする必要があるケースです。作者はMXFP8 Doryを例に挙げ、新しいフローでは予期しないキーがなく、生成結果もBF16の参照結果と一致したと報告しています。また、この修正では非圧縮チェックポイントに対して逆量子化を強制せず、量子化KVキャッシュのみを使うモデルに必要なスケールパラメーターを保持します。W8A8では、逆量子化を指定した場合に活性値の量子化を停止します。これらは現時点では作者の説明に基づく結果です。修正の詳細

デプロイを担当するエンジニアにとって注目点は、読み込み完了から初回推論までの処理の配分です。この設計から推測すると、展開が前倒しされれば、起動時間、初回リクエストのレイテンシ、読み込み後のメモリ使用量の測定結果の解釈が変わる可能性があります。ただし、提案には関連する性能数値が示されていないため、推論が高速化すると結論づけることはできません。

確認時点でPRはオープンのままレビュー待ちで、CIのコメントには逆量子化テストで新たな失敗が1件記載されています。その後、すべてのテストが通過したことも確認できていません。採用前にテスト状況を追跡し、手元のチェックポイントで名前のマッピング、スケールパラメーター、出力の一貫性を検証する必要があります。レビューとCIの状況