vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
作者: Congyue Deng
分类: cs.CV, cs.DC, cs.LG
发布日期: 2026-09-16
💡 一句话要点
提出vidax框架以解决TPU视频生成模型推理问题
🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)
关键词: 视频生成 JAX框架 TPU推理 时空模型 深度学习
📋 核心要点
- 现有的视频生成模型大多依赖于PyTorch/CUDA，无法充分利用Cloud TPU的资源，限制了其在长序列时空注意力任务中的应用。
- 本文提出vidax框架，提供一个JAX/Flax推理引擎和PyTorch到JAX的权重转换器，支持多种视频生成架构，消除了对PyTorch的依赖。
- 在TPU v4-8硬件上进行的基准测试显示，vidax在编译时间、延迟和峰值内存利用率上表现优异，解决了实际应用中的数值错误问题。
📝 摘要（中文）
现有的视频生成模型几乎仅以PyTorch/CUDA实现，导致Cloud TPU无法有效利用其大规模内存池进行推理。为此，本文提出vidax，一个开源的JAX/Flax推理引擎及零拷贝的PyTorch到JAX权重转换器，支持多种时空模型，包括扩散变换器和3D变分自编码器。vidax通过统一1D张量并行与DeepSpeed-Ulysses序列并行，集成TPU闪存注意力内核，并实现每层权重卸载，支持超出单设备内存的参考分辨率。我们在TPU v4-8硬件上进行了基准测试，记录了在检查点转换过程中出现的实际数值错误。vidax现已开源，作为JAX和TPU视频生成研究的基线。
🔬 方法详解
问题定义：现有的视频生成模型主要基于PyTorch/CUDA实现，导致Cloud TPU在推理时缺乏有效路径，无法充分利用其大规模内存池，限制了长序列时空任务的处理能力。
核心思路：本文提出vidax框架，旨在通过提供一个开源的JAX/Flax推理引擎和零拷贝的PyTorch到JAX权重转换器，消除对PyTorch的依赖，从而实现高效的视频生成。
技术框架：vidax框架整合了1D张量并行与DeepSpeed-Ulysses序列并行，采用单个JAX分片网格，集成TPU闪存注意力内核，并实现每层权重卸载，支持超出单设备内存的分辨率。
关键创新：vidax的主要创新在于其零拷贝权重转换机制和对多种时空模型的支持，显著提升了TPU在视频生成任务中的应用效率，与现有方法相比，提供了更高的灵活性和性能。
关键设计：在设计中，vidax采用了TPU闪存注意力内核，优化了内存管理，并通过每层权重卸载技术，确保在高分辨率下的有效内存利用，提升了整体推理性能。
🖼️ 关键图片
📊 实验亮点
在TPU v4-8硬件上的基准测试中，vidax展示了显著的性能提升，编译时间和延迟均优于现有的PyTorch实现，同时在峰值内存利用率上表现出色，解决了实际应用中的数值错误问题，为视频生成研究提供了新的基线。
🎯 应用场景
vidax框架的潜在应用场景包括视频生成、长序列时空数据处理以及多模态生成任务。其高效的推理能力和对TPU的优化利用，将推动相关领域的研究和应用，特别是在需要处理大规模视频数据的场景中，具有重要的实际价值和影响。
📄 摘要（原文）
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.