論文の概要: E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
- arxiv url: http://arxiv.org/abs/2609.35099v1
- Date: Mon, 28 Sep 2026 13:10:17 GMT
- ステータス: 翻訳完了
- システム内更新日: 2026-10-02 16:58:25.011578
- Title: E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
- Title(参考訳): E3J:GPUとTPU上でのユークリッド同変操作のための効率的かつオープンソースバックエンド
- Abstract要約: 我々は,GPUとTPUのJAXバインディングを備えた幾何学的深層学習アプリケーションのための,高速なユークリッド等価バックエンドe3jを提案する。
最適化されたPallasカーネルとアルゴリズムによる改善の両方を活用することで、このライブラリは最先端のスループットと、前方と後方の両方のパスでのランタイムを実現している。
- 参考スコア(独自算出の注目度): 1.232120906972724
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: We present e3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU. Leveraging both optimized CUDA and Pallas kernels and algorithmic improvements, the library achieves state-of-the-art throughput and runtime on both forward and backward paths. On a machine learning interatomic potential (MLIP) use case, it outperforms established backends, measuring up to 34% speed-up over cuEquivariance on water box NPT simulation using MACE, while remaining fully open source. E3j achieves over 80% efficiency over the H100 maximum memory bandwidth on tensor product operations, and in many cases more than doubles throughput of message passing convolutions forward compared to previously available backends. In addition, with the release of dedicated Pallas TPU kernel, e3j opens the possibility of large scale equivariant deep learning workloads on TPU architectures, which has so far been difficult to achieve. Our benchmarks show that e3j also achieves over 80% of a TPUv6e memory bandwidth, up to one order of magnitude more than e3nn-jax. The library is available on GitHub, PyPI and is released under an open source Apache 2.0 license.
- Abstract(参考訳): 本稿では,GPU と TPU の JAX バインディングを備えた幾何学的深層学習アプリケーションのための,高速なユークリッド等分散バックエンド e3j を提案する。
最適化されたCUDAとPallasのカーネルとアルゴリズムによる改善の両方を活用して、このライブラリは最先端のスループットと、前方と後方の両方のパスでのランタイムを実現する。
機械学習の原子間ポテンシャル(MLIP)のユースケースでは、確立されたバックエンドよりも優れており、完全なオープンソースを維持しながら、MACEを使用して、ウォーターボックスNTTシミュレーションにおけるcuEquivarianceの最大34%のスピードアップを計測する。
E3jはテンソル製品操作におけるH100最大メモリ帯域幅よりも80%以上の効率を実現しており、多くの場合、以前利用可能なバックエンドと比較してメッセージパッシング畳み込みのスループットが2倍以上になる。
さらに、専用のPallas TPUカーネルのリリースにより、e3jはTPUアーキテクチャ上での大規模な同変ディープラーニングワークロードの可能性を開放する。
我々のベンチマークでは、e3jはTPUv6eメモリ帯域の80%以上を達成しており、e3nn-jaxより1桁多い。
このライブラリはGitHub、PyPIで利用可能で、オープンソースのApache 2.0ライセンスでリリースされている。
関連論文リスト
- Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines [12.398594139609207]
本稿では,Google の Gemma 4 31B モデルを TPU ハードウェア上で実現した,ファインチューニングの最初のエンドツーエンドデモを紹介する。
Google TPU v5p-8のトレーニングにLoRAを使用し、推論にTPU v6e-8(Trillium)を使用すると、GPUネイティブなトレーニングレシピの移植に必要なコードレベルの適応の完全なセットを文書化する。
論文 参考訳(メタデータ) (2026-05-25T09:51:59Z) - CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation [51.72529978689561]
Agentは、カーネルの専門知識を3つのコンポーネントで開発する大規模なエージェント強化学習システムである。
AgentはKernelBench上で、トーチコンパイルよりも100%、100%、92%高速なレートを提供する。
論文 参考訳(メタデータ) (2026-02-27T18:58:05Z) - FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness [0.0]
FlashAttentionのようなメソッドは、不要なデータ転送を避けることによって、ネイティブのPyTorchよりもx6パフォーマンスの向上を実現している。
本稿では、ディープラーニングモデルのためのニューラルネットワークダイアグラムを拡張し、GPU階層間のリソース使用量とタスクの分散を検討する。
本研究では,中間レベルの擬似コードをダイアグラムで表現する手法を開発し,ハードウェア認識アルゴリズムをステップバイステップで導出できるようにする。
論文 参考訳(メタデータ) (2024-12-04T13:52:04Z) - Fully-fused Multi-Layer Perceptrons on Intel Data Center GPUs [3.7101665559244874]
本稿では,Intel Data Center GPU Max 1550用のMulti-formedLayer Perceptrons(MLP)の実装について述べる。
これにより算術強度が大幅に向上し,特に推論性能が向上することを示す。
論文 参考訳(メタデータ) (2024-03-26T11:38:39Z) - 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) - INR-Arch: A Dataflow Architecture and Compiler for Arbitrary-Order
Gradient Computations in Implicit Neural Representation Processing [66.00729477511219]
計算グラフとして表される関数を考えると、従来のアーキテクチャはn階勾配を効率的に計算する上で困難に直面している。
InR-Archは,n階勾配の計算グラフをハードウェア最適化データフローアーキテクチャに変換するフレームワークである。
1.8-4.8x と 1.5-3.6x の高速化を CPU と GPU のベースラインと比較した結果を示す。
論文 参考訳(メタデータ) (2023-08-11T04:24:39Z) - Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
この研究は、効率的でポータブルなDeep LearningとHigh Performance Computingカーネルを開発するためのフレームワークを導入している。
1)プロセッシングプリミティブ(TPP)を用いた計算コアの表現と,2)高レベルな宣言的手法でTPPのまわりの論理ループの表現の2つのステップでカーネルの開発を分解する。
我々は、スタンドアロンカーネルと、さまざまなCPUプラットフォームにおける最先端実装よりも優れたエンドツーエンドワークロードを使用して、このアプローチの有効性を実証する。
論文 参考訳(メタデータ) (2023-04-25T05:04:44Z) - Efficient and Generic 1D Dilated Convolution Layer for Deep Learning [52.899995651639436]
幅広いパラメータをカバーする汎用的な1D畳み込み層の効率的な実装を紹介します。
特にIntel AVX-512とAVX-512 BFloat16命令を含むアーキテクチャ向けに最適化されている。
本稿では,最適化された1次元畳み込み層の性能を,実際のゲノミクスデータセットを用いたエンドツーエンドニューラルネットワークトレーニングで実証する。
論文 参考訳(メタデータ) (2021-04-16T09:54:30Z) - Heterogeneous CPU+GPU Stochastic Gradient Descent Algorithms [1.3249453757295084]
ヘテロジニアスCPU+GPUアーキテクチャの深層学習のためのトレーニングアルゴリズムについて検討する。
私たちの2倍の目標 -- 収束率と資源利用を同時に最大化する -- は、この問題を難しくします。
これらのアルゴリズムの実装は,複数の実データセットよりも高速な収束と資源利用の両立を実現していることを示す。
論文 参考訳(メタデータ) (2020-04-19T05:21:20Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト(すべての情報・翻訳含む)の品質を保証せず、本サイト(すべての情報・翻訳含む)を使用して発生したあらゆる結果について一切の責任を負いません。