論文の概要: vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
- arxiv url: http://arxiv.org/abs/2609.18077v1
- Date: Wed, 16 Sep 2026 03:26:45 GMT
- ステータス: 翻訳完了
- システム内更新日: 2026-09-20 08:55:53.624101
- Title: vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
- Title(参考訳): vidax: アクセラレータメッシュ上のビデオ生成モデルのための統一JAXフレームワーク
- Abstract要約: オープンソースビデオ生成モデルは、ほとんどPyTorch/CUDAリファレンス実装として出荷される。
本稿では、JAX/Flax推論エンジンのvidaxと、ビデオ生成アーキテクチャのためのゼロコピーPyTorch-to-weightトランスレータを提案する。
- 参考スコア(独自算出の注目度): 8.113406913714298
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Open-source video generative models ship almost exclusively as PyTorch/CUDA reference implementations. This leaves Cloud TPU pods without a production-ready inference path, despite offering large, cost-effective accelerator memory pools ideal for long-sequence spatiotemporal attention. We present vidax, an open-source JAX/Flax inference engine and zero-copy PyTorch-to-JAX weight translator for modern video generation architectures. vidax covers a diverse set of spatiotemporal models --- including Diffusion Transformers, omnimodal Mixture-of-Transformers, 3D VAEs, text encoders, and native samplers --- with zero PyTorch dependency in the execution path. The framework unifies 1D tensor parallelism with DeepSpeed-Ulysses sequence parallelism on a single JAX sharding mesh, integrates TPU flash-attention kernels, and implements per-layer weight offloading to support reference resolutions that exceed single-device memory. We benchmark compile times, latency, and peak memory utilization on TPU v4-8 hardware, and document real-world numerical bugs surfaced during checkpoint translation. vidax is released open-source as a baseline for JAX and TPU video generation research.
- Abstract(参考訳): オープンソースのビデオ生成モデルは、ほとんどPyTorch/CUDAリファレンス実装として出荷される。
これにより、Cloud TPUポッドは、長時間の時空間的注意に最適な、コスト効率の大きいアクセラレータメモリプールを提供するにもかかわらず、プロダクション対応の推論パスを持たない。
本稿では、最新のビデオ生成アーキテクチャのためのオープンソースのJAX/Flax推論エンジンであるvidaxと、ゼロコピーのPyTorch-to-JAX重み変換器について述べる。
vidaxは、Diffusion Transformer、Omnimodal Mixture-of-Transformers、3D VAE、テキストエンコーダ、ネイティブサンプルなど、実行パスにおけるPyTorch依存性のない、さまざまな時空間モデルのセットをカバーする。
このフレームワークは、1つのJAXシャーディングメッシュ上で1Dテンソル並列性とDeepSpeed-Ulyssesシーケンス並列性を統一し、TPUフラッシュアテンションカーネルを統合し、単一デバイスメモリを超える参照解像度をサポートするために、レイヤ単位の重みオフロードを実装している。
我々は,TPU v4-8ハードウェア上でのコンパイル時間,レイテンシ,ピークメモリ利用率をベンチマークし,チェックポイント変換時に発生する実世界の数値バグを文書化する。
vidaxは、JAXおよびTPUビデオ生成研究のベースラインとして、オープンソースとしてリリースされた。
関連論文リスト
- Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent [10.241408048987038]
既存のGPUカーネル生成ベンチマークは、デプロイされたワークロードから分岐する合成またはキュレートされたソースから問題を引き出す。
Atrex-Benchは、30の演算子と440の形状が、計算限定のメモリリッチGPUのフルクラスタ生産推定トレースから直接サンプリングされるベンチマークである。
論文 参考訳(メタデータ) (2026-07-16T03:49:50Z) - Compiler-First State Space Duality and Portable $O(1)$ Autoregressive Caching for Inference [0.0]
Mamba-2のステートスペース双対性アルゴリズムは、XLAの融合とタイリングが実際に最適化したものにきれいにマッピングする。
我々は,XLA の標準プリミティブとして,完全な推論パス (プリフィル,キャッシュされた自己回帰復号化) を実装した。
この実装は、単一のJAXソースからCPU、NVIDIA GPU、Google Cloud TPUで無修正で実行される。
論文 参考訳(メタデータ) (2026-03-10T12:03:00Z) - pyFAST: A Modular PyTorch Framework for Time Series Modeling with Multi-source and Sparse Data [10.949140998070732]
pyFASTは時系列分析のための研究指向のPyTorchフレームワークである。
データエンジンは複雑なシナリオのために設計されており、マルチソースローディング、タンパク質配列処理、効率的なシーケンスおよびパッチレベルのパディング、動的正規化、マスクベースのモデリングをサポートする。
GitHubでMITライセンスでリリースされているpyFASTは、時系列の研究とアプリケーションを進めるための、コンパクトだが強力なプラットフォームを提供する。
論文 参考訳(メタデータ) (2025-08-26T10:05:47Z) - QuantCache: Adaptive Importance-Guided Quantization with Hierarchical Latent and Layer Caching for Video Generation [84.91431271257437]
Diffusion Transformers (DiTs) はビデオ生成において支配的なアーキテクチャとして登場した。
DiTには、計算コストやメモリコストの増大など、大きな欠点がある。
我々は,新しいトレーニングフリー推論アクセラレーションフレームワークQuantCacheを提案する。
論文 参考訳(メタデータ) (2025-03-09T10:31:51Z) - JORA: JAX Tensor-Parallel LoRA Library for Retrieval Augmented Fine-Tuning [16.86356520836045]
本稿では,Llama-2モデルのPEFT互換微調整のための新しいフレームワークについて紹介する。
我々のフレームワークは、JAXのジャスト・イン・タイム(JIT)コンパイルと、効率的なリソース管理のためにテンソルシャーディングを独自に利用しています。
実験では,Hugging Face/DeepSpeed実装を4GPUで実装するのに対して,GPUあたりのVRAMは半分以下であるのに対して,ランタイムでは12倍以上の改善が見られた。
論文 参考訳(メタデータ) (2024-03-17T23:02:04Z) - Green AI: A Preliminary Empirical Study on Energy Consumption in DL
Models Across Different Runtime Infrastructures [56.200335252600354]
トレーニング済みのモデルを、ネイティブな開発環境とは異なる環境にデプロイするのは、一般的なプラクティスです。
これにより、インフラを含むONNXや標準フォーマットとして機能するONNXなどの交換フォーマットが導入された。
論文 参考訳(メタデータ) (2024-02-21T09:18:44Z) - JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
我々は、GPU対応の効率と多くの一般的なMARL環境のサポートを組み合わせた、最初のオープンソースPythonベースのライブラリであるJaxMARLを紹介します。
我々の実験は、壁時計時間の観点から、JAXベースのトレーニングパイプラインが既存のアプローチの約14倍高速であることを示している。
また、人気の高いStarCraft Multi-Agent ChallengeのJAXベースの近似的な再実装であるSMAXを紹介し、ベンチマークする。
論文 参考訳(メタデータ) (2023-11-16T18:58:43Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト(すべての情報・翻訳含む)の品質を保証せず、本サイト(すべての情報・翻訳含む)を使用して発生したあらゆる結果について一切の責任を負いません。