vidax：加速器网格上视频生成模型的统一 JAX 框架

开源视频生成模型几乎完全作为 PyTorch/CUDA 参考实现提供。尽管提供了大型、经济高效的加速器内存池，非常适合长序列时空注意力，但这使得 Cloud TPU Pod 没有生产就绪的推理路径。我们推出了 vidax，一种开源 JAX/Flax 推理引擎和用于现代视频生成架构的零拷贝 PyTorch 到 JAX 权重转换器。 vidax 涵盖了多种时空模型——包括扩散Transformer、全模态混合Transformer、3D VAE、文本编码器和本机采样器——执行路径中对 PyTorch 的依赖性为零。该框架将 1D 张量并行性与 DeepSpeed-Ulysses 序列并行性统一在单个 JAX 分片网格上，集成 TPU 闪存注意力内核，并实现每层权重卸载以支持超过单设备内存的参考分辨率。我们对 TPU v4-8 硬件上的编译时间、延迟和峰值内存利用率进行基准测试，并记录检查点转换过程中出现的实际数字错误。 vidax 已开源发布，作为 JAX 和 TPU 视频生成研究的基准。