Jasmine让世界模型训练快10倍，支持从单机到百卡扩展。
Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase
- 基于JAX构建，代码改动少即可跨设备规模扩展。
- 复现CoinRun任务速度提升一个数量级，关键优化在数据加载与保存。
- 支持可复现训练，适合做模型架构对比实验的科研人员。
尽管世界模型被视为缓解机器人等领域数据稀缺问题的路径，但开放的训练基础设施仍不成熟。我们提出Jasmine，一个高性能的JAX框架世界模型代码库，仅需少量代码修改即可从单机扩展至数百加速器。相比以往开源实现，Jasmine在数据加载、训练和检查点保存方面进行性能优化，使CoinRun案例的复现速度提升一个数量级。该代码库保证完全可复现训练，并支持多种分片配置。结合精选的大规模数据集，为不同模型家族与架构消融实验提供了严谨的基准测试基础设施。
原文摘要 · Abstract (English)
While world models are increasingly positioned as a pathway to overcoming data scarcity in domains such as robotics, open training infrastructure for world modeling remains nascent. We introduce Jasmine, a performant JAX-based world modeling codebase that scales from single hosts to hundreds of accelerators with minimal code changes. Jasmine achieves an order-of-magnitude faster reproduction of the CoinRun case study compared to prior open implementations, enabled by performance optimizations across data loading, training and checkpointing. The codebase guarantees fully reproducible training and supports diverse sharding configurations. By pairing Jasmine with curated large-scale datasets, we establish infrastructure for rigorous benchmarking pipelines across model families and architectural ablations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。