AMDの研究者らはROCm Blogsで、ロボットポリシーのWorld Action Models(WAMs)であるFastWAMとDiT4DiTをAMD Instinct MI355X上で高速化する手法を公開しました。GEAKによるカーネル最適化と2ステップのデノイジングを組み合わせ、FastWAM/LIBEROのレイテンシを222.55 msから9.59 ms(約23倍)に短縮し、成功率97.00%を保っています。
AMD InstinctはAMDのデータセンターAI・HPC向けアクセラレーター製品群です。
何が新しいか
最適化は2段階です。第1段階は公開済みの重みを変えない学習不要の変更で、AMDのエージェント型GPUカーネル最適化フレームワークGEAKが見つけたカーネル最適化と、デノイジング回数の削減を組み合わせます。第2段階はアーキテクチャ変更で、アクションエキスパートを小型化するCompact Action DiTと、生アクションチャンクを連続潜在表現に圧縮するAction VAEを導入し再学習します。両者をまとめてCompact Action Expert(CAE)と呼んでいます。
MI355XでのFastWAMのプロファイルでは、Action DiTの10ステップのデノイジングループがエンドツーエンド推論時間の約87%を占めていました。GEAKの効果が大きかった項目は次の4つです。
- HIPグラフのキャプチャーと再生によるPython/HIP起動オーバーヘッドの削減
- Q・K・V投影の融合とクロスアテンション投影のグループ化による大きなGEMM化
- モジュール走査の重複排除やループ不変処理の巻き上げによるディスパッチ・データ移動の削減
- 単一フレーム向けにVAEを固定形状で特化し、3D畳み込みを等価な2D演算へ置換
GEAK適用後、Action DiTのデノイジングは約10倍、Video DiTのプリフィルは約4.5倍速くなりました。
数値(MI355X)
| 最適化 | FastWAM/LIBERO | DiT4DiT/LIBERO | DiT4DiT/RoboCasa |
|---|---|---|---|
| なし | 96.95% (222.55 ms) | 98.16% (81.47 ms) | 57.9% (82.5 ms) |
| GEAK | 96.86% (24.17 ms) | 98.24% (6.49 ms) | 58.2% (6.69 ms) |
| GEAK+2ステップ | 97.00% (9.59 ms) | 98.08% (5.65 ms) | 60.9% (5.35 ms) |
GEAK単体の高速化はFastWAMで9.2倍、DiT4DiT/LIBEROで12.6倍です。2ステップ設定はFastWAMが10→2、DiT4DiTが4→2ステップです。MI355Xではtorch.compileがeager PyTorch比で1.6〜2.6倍の短縮にとどまる一方、GEAKはtorch.compile比でFastWAMで3.7倍、DiT4DiTで7.7〜7.9倍高速でした。
CAEでは、Action VAEがFastWAMの32×7の生アクションチャンクを8×16の連続潜在表現に圧縮し、時間方向トークンを4分の1、スカラー数を43%減らします。FastWAM/LIBEROの成功率はCAEで97.28%(元は96.95%)、DiT4DiT/RoboCasaは56.8%(報告値56.7%)でした。CAEにGEAKと10→4ステップを組み合わせると、FastWAM/LIBEROは222.55 msから8.83 ms(約25倍)、成功率97.33%です。
LIBEROの成功率はMI355Xで11回の完全評価の平均、RoboCasaは単一評価実行の値です。GEAKの候補は、4スイートの7タスク×2エピソード、計14エピソードのスクリーニングで、3回の試行のいずれかが全14エピソードに成功した場合のみ合格とします。
試し方・実行環境
コードとチェックポイントはfaster-action-modelsリポジトリとHugging Faceで公開されています。
- ベースライン評価:
bash scripts/reproduce_baselines.sh - 学習不要最適化と
torch.compile比較の再現:bash scripts/reproduce_training_free.sh
MI355XではGPU物理エンジンとレイトレーサーで閉ループ全体をGPU上で実行します。Radeon PRO W7900ではMuJoCo物理演算をCPU、OpenGLレンダリングをGPUで実行します。最適化前の再現では、両GPUで公開値とLIBEROで0.65ポイント以内、RoboCasaで1.2ポイント以内の差でした。
背景
WAMは事前学習済みのビデオ/ワールドモデルで将来のシーン変化を予測し、その予測に基づいてアクションを生成します。FastWAMとDiT4DiTはVideo DiTとAction DiTを組み合わせる構造のため、制御ステップごとのVideo DiTプリフィルがレイテンシとメモリを消費し、ミリ秒単位の閉ループ制御で展開上の障害になると原文は説明しています。
まだコメントはありません。