MegaSlide-DiT:メモリ中心の適応と変形可能な局所注意機構による効率的なビデオ拡散
MegaSlide-DiTは、モデル状態をホストメモリに保持しGPUにストリーミングする方式と、3D変形可能スライド注意機構を組み合わせることで、105BパラメータのDiTを単一のH200 GPUで適応させるシステムレベルの手法を提案する。
拡散トランスフォーマー(DiT)に基づく高解像度ビデオ拡散モデルは、優れた画質を提供する一方で、その膨大なパラメータ数とメモリ要件が単一ワークステーションの能力をすぐに超えてしまいます。1000億パラメータを超えるDiTは、1テラバイト以上の永続状態(重み、最適化状態など)を必要とし、単純な時空間自己注意機構は系列長の二乗で計算量が増大します。これら2つの壁——パラメータメモリと活性化メモリ——により、研究者は大規模GPUクラスタなしでは巨大な生成モデルを適応させることが困難でした。
Jiacheng Liu氏とJason Liu氏の最新論文MegaSlide-DiTは、この問題をシステムの観点から再検討します。彼らは、1.5 TBのホストRAMを備えた単一のNVIDIA H200 GPU上で、事前学習済みの105B DiTを適応させるプロトタイプを提案します。その鍵となる洞察は、GPUがモデル状態をすべて保持する必要がないという点です。すべての永続的な重み、マスター重み、最適化器のモーメントはホストメモリに保持され、必要なときに一時的なシャードのみがオンデマンドでGPUにストリーミングされます。これにより、GPUメモリの制約を大幅に緩和できます。
同時に、研究者たちは二次のグローバル注意機構を、3D変形可能スライド注意機構(3D-DSA)に置き換えました。3D-DSAは、動きに適応した局所注意演算子で、ビデオフレーム内の時間的に変化する領域に焦点を当てることで、注意機構のメモリと計算の複雑さを系列長に対して線形に削減します。この設計により、メモリ使用量が大幅に節約されるとともに、計算速度も向上します。
論文では、詳細なメモリ使用量の内訳、実行トレース、評価結果が報告され、設計の有効性が実証されています。MegaSlide-DiTは、単一GPUで105Bモデルをゼロから訓練できると主張するものではなく、帯域幅の制限を魔法のように解決するものでもありません。むしろ、ハイエンドワークステーション上で大規模ビデオ拡散モデルの全パラメータ適応を可能にする実用的な道筋を提供します。この研究は、大規模AIモデルの導入障壁を低減する重要な一歩であり、特に予算制約のある研究チームや中小企業にとって大きな意味を持ちます。