Hugging Face pytorch_block_sparse リリース

Hugging Faceは、ブロック疎行列乗算を実装することで、より小さく高速なニューラルネットワークの作成を可能にする拡張機能 pytorch_block_sparse をリリースしました。このライブラリは、既存の疎代数計算における効率性のギャップに対処し、言語モデルのメモリ使用量の削減と生産コストの低減を実現します。

ブロック疎実装と使用方法

pytorch_block_sparse は、torch.nn.Linear のドロップイン・リプレイスメントとして機能する BlockSparseLinear モジュールを提供します。これにより、開発者は初期化時に密度パラメータを指定するだけで、最小限のコード変更でモデルに疎性を統合できます。

さらに、このライブラリには、既存のモデルを「オンザフライ」で修正できる BlockSparseModelPatcher が含まれています。これにより、元のモデルのソースコードを変更することなく、モデルを疎化し、その後通常通りにトレーニングを行うことが可能になります。

技術的基盤: NVIDIA CUTLASS

この拡張機能は、Yulhwa Kim による cutlass tilesparse の概念実証に基づいて構築されており、NVIDIA CUTLASS に基づくブロック疎行列乗算のために C++ CUDA テンプレートを利用しています。

CUTLASS は、高パフォーマンスな CUDA カーネルを実装するために使用される CUDA C++ テンプレートのコレクションであり、アセンブリ言語を必要とせずに cuBLAS に近いパフォーマンスレベルを実現します。このライブラリは、Ampere Tensor Core プリミティブと互換性があり、限定的な精度低下で 10 倍以上の高速化を実現できます。

パフォーマンスとメモリ効率

このライブラリにおける疎行列の現在のパフォーマンスは、cuBLAS で最適化された密行列の約 2 倍遅いですが、標準的な PyTorch の疎行列よりも大幅な改善されており、標準的な PyTorch の疎行列は密行列よりも 1 桁遅いものです。

パフォーマンスの向上は、疎性のレベルに応じてスケールします。例えば、75% の疎性を持つ行列は、密行列の約 2 倍高速です。メモリの節約効果はより顕著です。75% の疎性において、メモリ消費量は 4 倍削減されます。

今後の開発ロードマップ

pytorch_block_sparse の今後のイテレーションでは、初期化時に設定された固定の疎性パターンを超えた展開を目指しています。計画されている改善には以下が含まれます：

• Sparsity Pattern Optimization: 学習プロセス中に疎性パターンを最適化するために、パラメータの「有用性」を測定するツールを開発すること。
• Hardware Acceleration: さらなるパフォーマンス向上を実現するために、ブロック内に NVIDIA Ampere の 50% 疎性パターンを統合すること。
• CUTLASS Updates: 全体的な効率を高めるために、より新しいバージョンの CUTLASS を活用すること。