arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2610.20029

TrimSO(2): Fast and Memory-Efficient Training of Equivariant Interatomic Potentials

Chen Wang·Yuanchang Zhou·Guangming Tan·Mingzhen Li·Weile Jia
View original source ↗Read source document ↗Export metadata ↓

Abstract

Equivariant interatomic potentials enable accurate and efficient simulations of molecules and materials. Among them, models based on SO(2) convolutions reduce the cost of tensor products and achieve state-of-the-art accuracy, but their training remains time-consuming and memory-intensive. Existing training implementations materialize large edge tensors for representation transformations surrounding dense matrix multiplications, increasing memory usage and data movement. We observe that these transformations can be integrated into the model's existing dataflow, avoiding separate intermediate materialization. We introduce TrimSO(2), a fast and memory-efficient algorithm for SO(2)-based equivariant models. It realizes this principle through rotation on gather, recombination with activation, and inverse rotation on reduction, with support for the higher-order differentiation required for training. By avoiding intermediate materialization, TrimSO(2) delivers 2.71-16.73$\times$ operator speedups across forward, backward, and double backward and reduces model-level high-bandwidth memory traffic by 3.81-9.88$\times$. Finally, TrimSO(2) achieves 2.07-2.94$\times$ training speedups and 1.53-3.16$\times$ reductions in peak training memory over PyTorch implementations for official eSEN, UMA, and EquiformerV3 workloads on A100 and H100 GPUs.

Explore related subjects

cs.DC

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Chen Wang, Yuanchang Zhou, Guangming Tan, Mingzhen Li, Weile Jia. 2026-10-14. TrimSO(2): Fast and Memory-Efficient Training of Equivariant Interatomic Potentials. https://arxiv.org/abs/2610.20029

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

ZenFlow: Enabling Stall-Free Offloading for LLM Training

Fine-tuning large language models (LLMs) often exceeds GPU memory limits, prompting systems to offload model states to CPU memory. However, existing offloaded training frameworks like ZeRO-Offload treat all parameters equally and update the full model on the CPU, causing severe GPU stalls, where fast, expensive GPUs sit idle, waiting for slow CPU updates and limited-bandwidth PCIe transfers. We present ZenFlow, a new offloaded training framework that decouples updates between GPU and CPU by prioritizing important gradients. ZenFlow performs in-place updates of important gradients on GPU while offloading and accumulating less important ones on CPU without blocking training iterations, fully overlapping CPU work with GPU computation. To make this importance-aware design scalable, ZenFlow leverages a novel spatio-temporal locality of important gradients, concentrated in specific positions that remain stable across iterations, to make important gradient identification efficient and lightweight while preserving fidelity. Evaluation shows that ZenFlow achieves up to 5x end-to-end speedup, 2x lower PCIe traffic, 85% less GPU stalls, and up to 72% lower cloud GPU cost, all while preserving accuracy.

cs.DC↗
arXiv2026

TensorDex: A Compact, Tensor-Centric Storage System for Modern AI Models

Modern model hubs store hundreds of petabytes of large language models (LLMs), with fine-tuned variants dominating the storage footprint. These variants contain substantial cross-model redundancy that delta compression can exploit by storing only the difference between a target and a reference model. However, compression effectiveness depends critically on choosing a similar reference. At model-hub scale, this is challenging because model lineage metadata is often missing or unreliable, and different tensors within the same model may be most similar to tensors from different models. Consequently, model-level pairing leaves substantial redundancy unexploited. We present TensorDex, a lossless, tensor-centric storage system that performs delta compression at tensor granularity. TensorDex decomposes models into tensors, uses compact bit-level fingerprints to predict tensor-pair compressibility, and incrementally organizes tensors into multi-center clusters to select effective bases as new models arrive. Evaluated on 2,890 randomly sampled Hugging Face models, TensorDex reduces storage footprint by 70.5%, 37% lower than the state-of-the-art design. It achieves 22.9 GB/s compression and 28.4 GB/s decompression throughput, 3.86x and 1.49x faster than the next-best system, respectively.

cs.DC↗
arXiv2026

Coordination-Free Workload Regularization for Variable-Cost Distributed Training

Synchronous data-parallel training usually assigns every rank the same number of records, even when record costs vary by orders of magnitude. For variable-length models, this converts input variation into idle time because each step waits for its most expensive rank-local batch. We present Batch Cost Lookahead (BCL), a coordination-free data-loading policy that buffers a bounded rank-local window and repartitions it into fixed-cardinality packs with balanced estimated cost. BCL preserves the trainer's record-count contract and requires neither cross-rank metadata exchange nor runtime modification. In simulation over public workload distributions, enabling BCL reduces per-rank estimated batch-cost variability by up to 82.8%. Across 13 NeMo-RL SFT corpora on eight H200 GPUs, rank-local BCL runs 1.050 to 1.167 times faster than pooled MFFD. Compiled H200 measurements on an internal sequence workload show speedups of 1.197 to 1.343 times over packed execution from one to four nodes. We also identify a memory-feasibility boundary: a p99 cap for high-variance or OOM-prone data bounds outliers before BCL balances the rest. These results show that substantial variable-cost imbalance can be removed at the data-loader boundary, while clarifying when data preparation or broader coordination is still necessary.

cs.DC↗