論文の概要: DrJAX: Scalable and Differentiable MapReduce Primitives in JAX
- arxiv url: http://arxiv.org/abs/2403.07128v2
- Date: Wed, 17 Jul 2024 21:41:39 GMT
- ステータス: 処理完了
- システム内更新日: 2024-07-19 21:21:41.864986
- Title: DrJAX: Scalable and Differentiable MapReduce Primitives in JAX
- Title(参考訳): DrJAX: JAXにおけるスケーラブルで微分可能なMapReduceプリミティブ
- Authors: Keith Rush, Zachary Charles, Zachary Garrett, Sean Augenstein, Nicole Mitchell,
- Abstract要約: DrJAXは大規模分散および並列機械学習アルゴリズムをサポートするように設計されたライブラリである。
DrJAXはMapReduce計算のビルディングブロックをJAXのプリミティブとして組み込みます。
DrJAX計算はXLA HLOに直接変換することができ、幅広いMLトレーニングプラットフォームとの柔軟な統合を可能にする。
- 参考スコア(独自算出の注目度): 9.676195490442367
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: We present DrJAX, a JAX-based library designed to support large-scale distributed and parallel machine learning algorithms that use MapReduce-style operations. DrJAX leverages JAX's sharding mechanisms to enable native targeting of TPUs and state-of-the-art JAX runtimes, including Pathways. DrJAX embeds building blocks for MapReduce computations as primitives in JAX. This enables three key benefits. First, DrJAX computations can be translated directly to XLA HLO, enabling flexible integration with a wide array of ML training platforms. Second, DrJAX computations are fully differentiable. Last, DrJAX computations can be interpreted out to existing batch-processing compute systems, including traditional MapReduce systems like Apache Beam and cross-device compute systems like those powering federated learning applications. We show that DrJAX provides an easily programmable, performant, and scalable framework for parallelized algorithm development. DrJAX is available at \url{https://github.com/google-research/google-research/tree/master/drjax}.
- Abstract(参考訳): 我々は、MapReduceスタイルの操作を使用する大規模分散および並列機械学習アルゴリズムをサポートするように設計されたJAXベースのライブラリであるDrJAXを紹介する。
DrJAXは、JAXのシャーディングメカニズムを活用して、TPUとPathwaysを含む最先端のJAXランタイムのネイティブターゲティングを可能にします。
DrJAXはMapReduce計算のビルディングブロックをJAXのプリミティブとして組み込みます。
これにより、3つの重要なメリットが実現される。
まず、DrJAX計算をXLA HLOに直接変換することで、幅広いMLトレーニングプラットフォームとの柔軟な統合を可能にします。
第二に、DrJAX計算は完全に微分可能である。
最後に、DrJAX計算は、Apache Beamのような従来のMapReduceシステムや、フェデレートされた学習アプリケーションを動かすようなクロスデバイスコンピューティングシステムを含む、既存のバッチ処理コンピューティングシステムに解釈することができる。
DJAXは並列化アルゴリズム開発のための、プログラムが容易で、パフォーマンスが高く、スケーラブルなフレームワークを提供する。
DrJAX は \url{https://github.com/google-research/google-research/tree/master/drjax} で入手できる。
関連論文リスト
- Benchmarking Predictive Coding Networks -- Made Simple [48.652114040426625]
まず,性能と簡易性を重視したPCXというライブラリを提案する。
私たちはPCXを使って、コミュニティが実験に使用する大規模なベンチマークを実装しています。
論文 参考訳(メタデータ) (2024-07-01T10:33:44Z) - JaxDecompiler: Redefining Gradient-Informed Software Design [0.0]
JaxDecompilerは、任意のJAX関数を編集可能なPythonコードに変換するツールである。
この記事では、JAX関数を編集可能なPythonコードに変換するツールであるJaxDecompilerを紹介します。
論文 参考訳(メタデータ) (2024-03-14T20:32:31Z) - BlackJAX: Composable Bayesian inference in JAX [8.834500692867671]
BlackJAXはサンプリングと変分推論アルゴリズムを実装するライブラリである。
Pythonで書かれており、JAXを使ってNumpPyのようなサンプルをコンパイルし、CPU、GPU、TPU上で変分メソッドを実行する。
論文 参考訳(メタデータ) (2024-02-16T16:21:02Z) - 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) - JaxPruner: A concise library for sparsity research [46.153423603424]
JaxPrunerはスパースニューラルネットワーク研究のためのオープンソースライブラリである。
メモリとレイテンシのオーバーヘッドを最小限に抑えながら、一般的なプルーニングとスパーストレーニングアルゴリズムを実装している。
論文 参考訳(メタデータ) (2023-04-27T10:45:30Z) - SequeL: A Continual Learning Library in PyTorch and JAX [50.33956216274694]
SequeLは継続学習のためのライブラリで、PyTorchとJAXフレームワークの両方をサポートする。
それは、正規化ベースのアプローチ、リプレイベースのアプローチ、ハイブリッドアプローチを含む、幅広い連続学習アルゴリズムのための統一インターフェースを提供する。
私たちはSequeLをオープンソースライブラリとしてリリースし、研究者や開発者が自身の目的で簡単にライブラリを実験し拡張することができます。
論文 参考訳(メタデータ) (2023-04-21T10:00:22Z) - Extending Compositional Attention Networks for Social Reasoning in
Videos [84.12658971655253]
ビデオにおけるソーシャルインタラクションを推論するタスクのための,新しいディープアーキテクチャを提案する。
構成注意ネットワーク(MAC)の多段階推論機能を活用し,マルチモーダル拡張(MAC-X)を提案する。
論文 参考訳(メタデータ) (2022-10-03T19:03:01Z) - SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive
Knowledge Graphs [147.73127662757335]
我々は、知識グラフ(KG)におけるシングルホップおよびマルチホップ推論のための最初の汎用フレームワークであるスケーラブルなマルチホップ推論(SMORE)を提案する。
シングルマシンのSMOREはFreebase KG(86Mエンティティ、338Mエッジ)でマルチホップ推論を行うことができる。
SMOREは、従来のマルチホップKGフレームワークよりもスループット(トレーニング速度)を、最小のGPUメモリ要件で2.2倍向上させる。
論文 参考訳(メタデータ) (2021-10-28T05:02:33Z) - SymJAX: symbolic CPU/GPU/TPU programming [9.868558660605995]
SymJAXは、グラフ入力/出力/更新を単純化し、一般的な機械学習およびディープラーニングアプリケーションのための追加機能を提供するJAXのシンボリックプログラミングバージョンである。
ユーザの視点からは、SymJAXはLasagneのようなディープラーニング機能とともに、高速なグラフ最適化/コンパイルと幅広いハードウェアサポートを備えたLa Theanoエクスペリエンスを提供します。
論文 参考訳(メタデータ) (2020-05-21T13:37:25Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト(すべての情報・翻訳含む)の品質を保証せず、本サイト(すべての情報・翻訳含む)を使用して発生したあらゆる結果について一切の責任を負いません。