TT-NNモデル実装 実践編

Transcript

1. Tenstorrent Tech Talk #5 TTNNモデル実装 実践編 ~PytorchモデルのTTNN移植から最適化まで~
2. Tenstorrent OSS Stack Low level library/Kernel DNN Library / Python API DNN Compiler
3. Tenstorrent OSS Stack Low level library/Kernel DNN Library / Python API このSessionでのテーマ DNN Compiler
4. Part 1: Pytorch → TTNN Porting
5. What is TTNN ? • TT-MetaliumのNeural Network用のライブラリ • PytorchライクなNN命令が実装されており、C++とPythonで実装可能 • この命令を使うことでDNNモデルをTTハードウェアに対して実装することができる • ML FrameworkからTTハードウェアをターゲットにする場合 • Pytorch, JaxといったFrameworkでTTハードウェアをターゲットにして計算をする場合は、BackendでTTNNの 命令に変換する必要がある • そのためのDNNコンパイラやBackend実装はTenstorrentで提供している • Forge Compiler: Supports Pytorch/Jax/ONNX/TF/PaddlePaddle • PyTorch 2.0 TT-NN Backend
6. Key Features of TTNN • 高レベルのNeural Network Operations • 行列積, 畳み込み, Attention, データ移動, 集団通信(CCL), element-wise Ops, Reduction, 損失関 数, Pooling, などの主要なNN命令を最適化実装 • 柔軟なTensor Library • Tensor APIを通じて、TTアクセラレータに対するデータレイアウトをコントロールすることができ る • Naitve Multi-Device Support • デバイスメッシュ構成をサポートし、デバイスクラスタ全体でのスケーリングを容易に実現 • 計算グラフの可視化・Profilerツールの提供 • https://github.com/tenstorrent/ttnn-visualizer
7. How to use TTNN Install TTNN library from PyPI Simple TTNN Example
8. New model bringup flow in TTNN Reference Torch model Gen model graph Extract model summary • Operations • Parameters • Shape/Size Convert Torch to TTNN Ops Per Op PCC Impl Module Per Module PCC End-to-End End-to-End PCC Optimization Step by Step Implementation Mapping Torch to TTNN Ops Request to TT- Metal Develop Ops/Features in TT-Metal If no If yes
9. Model graph and summary TTNNへ移植するモデルのPytorch実装を解析する • Example Using: torchsummary, torchviz
10. Mapping Torch to TTNN Ops https://docs.tenstorrent.com/tt-metal/latest/ttnn/ttnn/api.html Yolov4のケース: • Conv2d → ttnn.conv2d • BatchNorm2d → ttnn.batch_norm • ReLU → ttnn.relu / ttnn.leaky_relu • MaxPool2d → ttnn.max_pool2d • Concat → ttnn.concat • Add → ttnn.add • Upsample → ttnn.upsample
11. Write TTNN model from Torch
12. Experimental: Forge compiler can generate TTNN C++ Code Forge frontend TT-MLIR TTNN/TT-Metalium Runtime TTNN Python /C++ EmitPy/C Forge compiler Stack Forge lowering
13. Part 2: TTNN Performance 最適化
14. TT-CNN • Vision model用のTTNNライブラリ • https://github.com/tenstorrent/tt-metal/tree/main/models/tt_cnn • チューニングされたAPIとパイプラインが提供されており、Vision系モデルの実装を容易にする TT-CNN Pipeline API • Traced Execution: Dispatchコマンドを記録し、実 行時に再生することでホストのOverheadを削減 • Multi-command Queue: 並列I/Oのキュー
15. TT-CNN: 4 level Executor Level 1: Model Executor (Base) • ケース: 初期実装, デバッグ, 小規模モデル • 特徴: Single Command queue, Tracingなし Level 2: Traced Model Executor (Trace最適化) • ケース: Host オーバーヘッドの削減が必要な場合 • 特徴: Metal Traceによる命令の事前記録と再生 Level 3: MultiCQTracedModelOverlappedInputExecutor (I/O並列化) • ケース: 入力転送のOverlapが必要な場合 • 特徴: 2つのCommand Queueで転送と計算を並列化（with Trace最適化） • CQ0: モデル実行・出力読み出し、CQ1: 入力書き込み Level 4: MultiCQTracedModelPipelinedIOExecutor (完全パイプライン化) • ケース: 最高性能の最適化 • 特徴: 入出力転送を完全に分離（with Trace最適化） • CQ0: モデル実行のみ、CQ1: 全ての入出力転送
16. Metal Trace • 実行する命令群のDispatch commandをDRAM bufferに記録して、実行時にそれらを再生する ことで、モデル実行時のHostオーバーヘッドを削減する • Dispatchにかかるホスト側の時間が、デバイス側の実行時間を超えたり、ホスト側の時間が大きく占める場合に 効果を発揮する Host boundな処理: デバイス上の計算間で隙間が多い Metal Traceで改善したケース
17. Metal Trace 1.コマンドの記録 • `begin_trace_capture`と`end_trace_capture` の間で実行される操作のコマンドがDRAM bufferに保存される • ※ `trace_region_size` を指定して、Trace bufferのためのDRAM領域のサイズを決定する必要がある 2.静的パラメータ • 全ての操作パラメータ（入出力TensorのShape, Addressなど）が静的に保存される • 制約: つまり、生成モデルといったシーケンス長が変化して静的に決定できないモデルに対して利用できない 3.デバイス上での再生: 記録されたコマンドはHostの介入なしにデバイス上で直接再生される
18. Metal Trace: Sample code
19. Metal Trace: deallocation • ttnn.deallocate でメモリを解放できる • Trace capture前には不要なTensorを解放する • モデルの出力を強制解放することで、L1メモリが解放されて、 Trace用の永続Bufferが適切なアドレスに配置される • トレース中、永続バッファは解放しない • 中間Tensorは演算後に即解放することが原則
20. Multiple Command Queues • TTデバイスにおける、I/O転送と計算実行を Overlapさせることでスループットを向上させ る • Metalでは最大2つの独立したCommand Queueをサポートしており、これによってデバ イスへのコマンドを並列してDispatchすること ができる • 一般的な構成では、CQ0が計算実行、CQ1がデータ転送 を担当するが、両方のCQは全てのDispatchコマンドを サポートするので、ユーザによって柔軟に役割を変える ことができる • それぞれのCQは互いに独立しているため、入力 書き込み用と処理実行用のQueueの命令順序を 保証するには、Queue間の同期にイベントを使 用する必要がある
21. Multiple Command Queues: Ops on CQ0, Read/Write on CQ1 • 計算処理を行うCQ0と入出力を行うCQ1を分離することで「演算時間が長い → I/Oを並列実行 して全体を高速化（パイプライン並列化）」とする • 各Queue間で依存関係を制御する4つのイベント • 1. 入力の書き込み完了 → CQ0が安全に演算を開始できるようにする • 2. CQ0が入力の使用完了 → CQ1が次の入力を上書きできるようにする • 3. CQ0が演算完了 → CQ1が出力読み取りできるようにする • 4. CQ1が出力読み取り完了 → CQ0が出力領域を上書きできるようにする
22. Multiple Command Queues: Sample code Ops and Output Readback on CQ0, Input Writes on CQ1
23. TT-CNN Level 4: Full Pipeline Full Pipelineのメモリフロー Input: [CQ1] Host → DRAM Execute: [CQ0] DRAM → L1 → Compute → L1 → DRAM Output: [CQ1] DRAM → Host Timeline （以下、Tの中で並列実行） T0: [CQ1: Input0 Write] T1: [CQ0: Input0 Execute] [CQ1: Input1 Write] T2: [CQ0: Input1 Execute] [CQ1: Output0 Read + Input2 Write] T3: [CQ0: Input2 Execute] [CQ1: Output1 Read + Input3 Write] …
24. Fast Dispatch • TT-Metalスタックにおける非同期Command Queueの仕組みで、ホストCPUのオーバーヘッドを削減し、 デバイスの利用率を最大化するための重要な機能 • Fast Dispatchは各Command Queueごとに1つのRISC-V Coreを割り当てて、非同期Commandキュー イングを担当する • これにより、ホストからデバイス操作の管理を切り離すことができ、デバイスのUtilizationが上がる • TT_METAL_SLOW_DISPATCH_MODE=1 でSlow Dispatchに切り替えられる（デバッグ用）
25. Metal Trace and Multiple Command Queues Together • Metal TraceとMultiple Command Queuesの最適化を組み合わせることで、高い性能向上が実現できる • HostがDeviceを大きく先行して動作し、多くの後続の反復処理を事前にQueueに入れる一方で、Device 上の命令は継続的に実行されるため、命令間の遅延はほとんどないか、ゼロになる。 例: Deviceが1回の反復処理を完了する前にHost側で10回の反復処理をQueueに入れており、命令間の隙間がほとんどない
26. How to use Tracy in Tensix • https://docs.tenstorrent.com/tt-metal/v0.55.0/tt-metalium/tools/tracy_profiler.html • TT forkのTracyをインストールして、起動してGUIを立ち上げる • TT-MetalでTracy scriptを経由してDNNモデルの計算を実行 • TT_METAL_DEVICE_PROFILER は、Tensix Baby RISC-V Cores上のProfilingをONにするために 必要な環境変数（DefaultではOFF） • Baby Coresによる制御処理とKernel Dispatch tasksがTraceされる（※ Tensix Engine上の具体的 な計算まではcaptureされないため、ttnn-visualizerを利用する） 起動
27. Memory on TTNN Interleaved layout Shared layout ラウンドロビンでDRAM Banksへ均等にTensorのPageを保存 シンプルな構成なのでアクセスが簡単だが、 遠いDRAMへのアクセスが必要だったり、Unusedな領域が生じる コアに近いDRAMへTensorを配置することができるので高速 一方で、Interleaved Layoutよりも制御が複雑 TTNN Memory配置は、DRAMとL1のMemory Locationと、InterleavedとSharedのMemory Layoutを組み合わせる
28. TTNN Memory: Interleaved • DRAM Interleaved • 大容量、任意のサイズのTensor配置 • 入力の大きいTensor, パラメータの保存 • L1よりも低速転送、ラウンドロビン • L1 Interleaved • L1転送の広帯域、Shardingより設定が簡単 • Shardが不要な小規模演算、単純計算の中間Bufferで利 用 • L1xコア分なので大きいTensorには不向き
29. TTNN Memory: Sharding • DRAM Sharding • DRAMの容量とShardingの並列性を両立 • L1へのreshardが高速 • Shard次元はコア数で割り切れる必要がある • L1 Sharding • 並列計算の最大化、L1帯域幅 x コア数 • Reshard Free (演算間でShardingが一致すればコピー不要 • L1xコア分なので大きいTensorには不向き
30. DataFormat and Math Fidelity on TTNN Block Floating Point Math Fidelity GroupingしてExponentを共通化することで圧縮したData Format 5bit x 7bitの乗算器が1Cycleで計算できる それ以上のサイズは最大4Stepsに分けて計算することで精度を担保 1~2Stepsで計算をストップさせることで高速化が可能
31. Math Fidelity Level
32. LLM Configuration: Performance mode • Example: Llama-3.1-70B • ベースはBFP8/HIFI2 • Performance/Accuracyのバランス重視 • ActivationはInput依存 (多くはBF16) • PrefillのSDPAはHIFI4 • 後続のKV cacheに影響を与えるため精度が必要 • Performance Mode • FF1_FF3 (MLP): BFP4/LOFI • MLPの低精度化は計算量・帯域幅の削減効果が高い • BFPの低精度に加えて、LOFIによって1 Clockで計算 できるように設定して高速化 tt-transformers model_config.py Performance mode config
33. fp32_dest_acc_en and packer_l1_acc • Fp32_dest_acc_en: Dstレジスタのデータ幅を制御し、累積精度を決定 • FP32で保存するためAccumulatorで高精度にデータを制御できる • 速度低下、レジスタ数半減 • Packer_l1_acc: PackerによるL1メモリ内での累積を有効化し、中間結果をL1に保持して再利用 • L1で累積することで、DRAM転送ボトルネック削減 • L1サイズの制約により、大きいTensorに対しては利用不可（OoMが発生） stable_diffusion_xl_base configuration
34. TTNN Sharding • Tensorを複数のコアに分割して並列処理 & 高帯域幅で高速 • 4つのSharding戦略: 戦略 分割方向 主な用途 AutoSharded 自動選択 Prototype, 自動最適化 HeightSharded 高さ(H)方向 バッチ処理、Matmul WidthSharded 幅(W)方向 Channel並列、DRAM I/O BlockSharded 二次元 (HxW) 方向 Conv2D （最大並列度） HeightSharded BlockSharded
35. Sharded Input Data Movement • Output = Activation * Weights を例に取る • 単一の出力 C[I, j] を計算するには、 • Activationの完全な行 A[i, :] が必要 • Weightsの完全な列 B[:, j] が必要 Height Sharding Width Sharding Block Sharding 各コアが完全な行を保持 Weightsは全コアに複製 独立して計算可能 コア間データ移動なし 全コアが最終的に同じ完全な出力を持つ チャネル並列処理 全コア間でデータ共有が必要 Broadcast回数 = コア数 通信オーバヘッドが大きい Width Shardingより通信量が少ない 行間は完全並列 2次元並列処理では最大スループット 行内でのBroadcastが必要 Width Shardingよりは少ないが通信アリ
36. Sharding Strategy: Vision Model • Conv2D • バッチサイズが大きい: HEIGHT_SHARDED • 空間サイズ(HxW)が大きい: BLOCK_SHARDED • 推論(Batch=1): BLOCK_SHARDED • Fully Connected • WIDTH_SHARDED • FC層は[Batch, Features] であり、出力チャネル方向で 分割して並列化する(Channel 並列) • Average Pooling • WIDTH_SHARDED • 各チャネルで独立したHxWのReductionがあるため、 [Batch, Channel]ごとに並列計算
37. TT-CNN Optimization Summary Optimization Method | Performance • Metal Trace | +++++ • 2CQ Pipeline | +++++ • Double Buffering | +++ • Memory Sharding | +++ • Sharding Strategy | +++++ • Channel Slicing | ++++ • Deallocation | ++ • Config in DRAM | + • Data Format | ++++ • Math Fidelity | ++++ • Haloing | ++++ • Reshard Optimization | ++ 並列実行・ パイプライン化 Memory最適化 Operation最適化
38. Conclusion • TT Device上でDNNモデルを動かしたい場合は、TTNNで実装する • PytorchのモデルからTTNNへ移植するのが最も簡単 • 対応する命令を置き換えていくイメージ • TTNNへの移植が済んだら、最適化を行う • Host TT Deviceのやり取りを削減（パイプライン化） • Data Format/Math Fidelityを適切に設定（精度とトレードオフ） • Sharding戦略で、コアにTensorを分散させて並列計算 • Etc.. • さらなる最適化はTT-Metal/LLK Kernelレベルで可能 • ttnn-visualizerでProfilingをして、最適化のための解析をする