論文の概要: Orbax: Distributed Checkpointing with JAX
- arxiv url: http://arxiv.org/abs/2605.23066v2
- Date: Tue, 26 May 2026 19:23:57 GMT
- ステータス: 翻訳完了
- システム内更新日: 2026-05-28 17:38:54.877696
- Title: Orbax: Distributed Checkpointing with JAX
- Title(参考訳): Orbax: JAXによる分散チェックポイント
- Authors: Colin Gaffney, Shutong Li, Daniel Ng, Anastasia Petrushkina, Niket Kumar, Adam Cogdell, Mridul Sahu, Yaning Liang, Nikhil Bansal, Justin Pan, Angel Mau, Abhishek Agrawal, Marco Berlot, Ruoxin Sang, Kiranbir Sodhia, Rakesh Iyer,
- Abstract要約: Orbaxは、分散アクセラレータシステムの複雑さを抽象化するJAXネイティブなチェックポイントライブラリである。
我々は、PyTorchの競合よりパフォーマンスが3.5$times$を節約し、2$times$をロードすることを示した。
- 参考スコア(独自算出の注目度): 2.2139026464835574
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: In a landscape of high-performance distributed ML systems, JAX has emerged as a framework of choice. However, JAX's modular design philosophy leaves it without a standardized checkpointing solution. In this paper, we introduce Orbax, a modular, JAX-native checkpointing library that abstracts the complexities of distributed accelerator systems while also providing flexibility for user-friendly checkpoint manipulations throughout the ML model lifecycle. We demonstrate performance exceeding comparable PyTorch competitors by up to 3.5$\times$ for saving and 2$\times$ for loading. The library is available at https://github.com/google/orbax.
- Abstract(参考訳): 高性能分散MLシステムの展望の中で、JAXは選択のフレームワークとして現れました。
しかし、JAX のモジュラー設計哲学は、標準化されたチェックポイントソリューションなしでそれを残します。
本稿では,分散アクセラレータシステムの複雑さを抽象化し,MLモデルライフサイクルを通じてユーザフレンドリなチェックポイント操作の柔軟性を提供する,モジュール型のJAXネイティブなチェックポイントライブラリであるOrbaxを紹介する。
我々は、PyTorchの競合よりパフォーマンスが3.5$\times$を節約し、2$\times$をロードします。
ライブラリはhttps://github.com/google/orbax.comから入手できる。
関連論文リスト
- MASEval: Extending Multi-Agent Evaluation from Models to Systems [77.09381093359598]
MASEvalはフレームワークに依存しないライブラリで、システム全体を分析の単位として扱う。
研究者はエージェントシステムのすべてのコンポーネントを探索し、原則化されたシステム設計のための新たな道を開くことができる。
論文 参考訳(メタデータ) (2026-03-09T18:46:17Z) - jinns: a JAX Library for Physics-Informed Neural Networks [0.7366405857677227]
jinnsは物理インフォームドニューラルネットワークのためのオープンソースのPythonライブラリである。
JAXエコシステムに根ざしたこのフレームワークは、リアルプロブレムを効率的にプロトタイピングするための多用途フレームワークを提供する。
論文 参考訳(メタデータ) (2024-12-18T18:21:41Z) - DrJAX: Scalable and Differentiable MapReduce Primitives in JAX [9.676195490442367]
DrJAXは大規模分散および並列機械学習アルゴリズムをサポートするように設計されたライブラリである。
DrJAXはMapReduce計算のビルディングブロックをJAXのプリミティブとして組み込みます。
DrJAX計算はXLA HLOに直接変換することができ、幅広いMLトレーニングプラットフォームとの柔軟な統合を可能にする。
論文 参考訳(メタデータ) (2024-03-11T19:51:01Z) - BlackJAX: Composable Bayesian inference in JAX [8.834500692867671]
BlackJAXはサンプリングと変分推論アルゴリズムを実装するライブラリである。
Pythonで書かれており、JAXを使ってNumpPyのようなサンプルをコンパイルし、CPU、GPU、TPU上で変分メソッドを実行する。
論文 参考訳(メタデータ) (2024-02-16T16:21:02Z) - XLB: A differentiable massively parallel lattice Boltzmann library in Python [0.0]
JAX プラットフォームをベースとした Python ベースの差別化可能な LBM ライブラリである XLB ライブラリを紹介する。
XLBの差別化性とデータ構造は、広範囲にわたるJAXベースの機械学習エコシステムと互換性がある。
XLBは数十億のセルでのシミュレーション処理に成功し、毎秒ギガスケールの格子更新を実現している。
論文 参考訳(メタデータ) (2023-11-27T18:50:37Z) - 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) - SequeL: A Continual Learning Library in PyTorch and JAX [50.33956216274694]
SequeLは継続学習のためのライブラリで、PyTorchとJAXフレームワークの両方をサポートする。
それは、正規化ベースのアプローチ、リプレイベースのアプローチ、ハイブリッドアプローチを含む、幅広い連続学習アルゴリズムのための統一インターフェースを提供する。
私たちはSequeLをオープンソースライブラリとしてリリースし、研究者や開発者が自身の目的で簡単にライブラリを実験し拡張することができます。
論文 参考訳(メタデータ) (2023-04-21T10:00:22Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト(すべての情報・翻訳含む)の品質を保証せず、本サイト(すべての情報・翻訳含む)を使用して発生したあらゆる結果について一切の責任を負いません。