sophgo/tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

解決する課題

TPU-MLIRは、さまざまな主流フレームワークの学習済みニューラルネットワークをTPUハードウェアにデプロイする問題を解決します。高レベルのモデル定義（PyTorchやONNXなど）と、TPUチップ上での効率的な実行に必要な特定のバイナリ形式（bmodel）との間のギャップを埋めます。

仕組み

このプロジェクトは、MLIR（Multi-Level Intermediate Representation）に基づいたコンパイルパイプラインを使用して、モデルを以下のいくつかの段階を経て変換します：

  1. フロントエンド・インポート: PyTorch、ONNX、TFLite、Caffe、またはHuggingFaceのモデルがインポートされ、統一されたMLIR表現に変換されます。
  2. ローリング（Lowering）: コンパイラは、パターン書き換えとメモリプランニングを使用して、高レベルの表現をTPU固有のダイアレクト（TopおよびTpu）に変換します。
  3. 量子化: さまざまな量子化技術（F32、BF16、F16、INT8）とキャリブレーションを適用し、ハードウェア効率のためにモデルを最適化します。
  4. デプロイ: 最終的な出力は、ターゲットのTPUプロセッサ上で実行可能なbmodelファイルです。

対象者

ディープラーニングモデル（大規模言語モデルおよびビジョンモデルを含む）をSophgo TPUハードウェアにデプロイする必要があるAIエンジニアおよび開発者向けに設計されています。

ハイライト

  • 幅広いフレームワークサポート: PyTorch、ONNX、TFLite、Caffe、HuggingFaceに対応。
  • LLM最適化: HuggingFace LLM向けの専用ワンショット変換。履歴KVキャッシュおよびチャンク推論をサポート。
  • 包括的な量子化: 対称/非対称INT8、F16、BF16をサポートし、AWQ、GPTQ、AutoRoundのパススルーにも対応。
  • プロダクション・ツール: 汎用推論ランナー（model_runner）、bmodel検査ツール（model_tool）、および精度検証ユーティリティが含まれます。