TorchSim: Batched MLIP Simulation Engine

• TorchSim is an open-source PyTorch-native simulation engine for MLIPs that enables batched, GPU-accelerated atomistic simulations.
• It integrates both classical and machine-learned interatomic potentials within a unified tensor framework to optimize throughput and memory management.
• Its innovative AutoBatching and mixed-size execution model yield significant speedups over traditional single-system simulation methods.

Searching arXiv for TorchSim and closely related atomistic simulation frameworks to ground the article in current papers. TorchSim is an open-source atomistic simulation engine implemented in PyTorch and designed for the Machine Learned Interatomic Potential (MLIP) era. It rewrites core atomistic simulation primitives in PyTorch so that simulation, force-field inference, and autograd operate in a common tensor framework, with an execution model centered on batched simulation of multiple systems on modern GPUs rather than the conventional single-system workflow. In the published description, TorchSim supports molecular dynamics integrators, structural relaxation optimizers, machine-learned and classical interatomic potentials, batching with automatic memory management, differentiable simulation, and interoperability with materials informatics tools (Cohen et al., 8 Aug 2025).

1. Concept and design scope

TorchSim was introduced to address a mismatch between contemporary MLIP workflows and the design assumptions of traditional atomistic simulation packages. The motivating claim is that widely used packages such as ASE, LAMMPS, and HOOMD-blue remain valuable, but were not designed around the dominant MLIP pattern in which models are trained in PyTorch, executed on GPUs, and evaluated repeatedly across many structures (Cohen et al., 8 Aug 2025). TorchSim therefore places PyTorch at the center of the simulation stack rather than treating it as an external inference backend.

Its stated scope is broader than molecular dynamics alone. The engine supports molecular dynamics integrators, structural relaxation optimizers, both machine-learned and classical interatomic potentials, differentiable simulation, and integration with common materials informatics workflows (Cohen et al., 8 Aug 2025). The same source describes TorchSim as a practical alternative to conventional packages for workloads in which MLIP inference dominates runtime and GPU utilization is the limiting systems concern rather than the availability of classical simulation kernels.

A defining distinction is execution granularity. Conventional engines generally simulate one system at a time, whereas TorchSim advances multiple unrelated systems concurrently in a batched representation (Cohen et al., 8 Aug 2025). This is presented not as a convenience feature but as the central systems-level choice through which accelerator utilization is increased. A plausible implication is that TorchSim should be understood as a simulation engine optimized for throughput under heterogeneous small- and medium-sized workloads, rather than as a direct reimplementation of legacy high-performance molecular dynamics codes.

2. PyTorch-native simulation model

The core technical move in TorchSim is the reimplementation of atomistic simulation primitives in PyTorch. According to the paper, these primitives include force evaluation through MLIPs or classical potentials, integrator operations, structural relaxation steps, and state transformations (Cohen et al., 8 Aug 2025). Because these operations are expressed as PyTorch tensor programs, TorchSim inherits GPU execution, autograd-based differentiability, and compatibility with the PyTorch model ecosystem.

The model interface is standardized through torch_sim.models.ModelInterface, which specifies the expected input-output schema for energies, forces, and stress (Cohen et al., 8 Aug 2025). The same description notes that the interface is intended to remain extensible to future outputs such as Born charges and polarization, spins or magnetic moments, and partial charges. This suggests that TorchSim is structured around a simulation-state abstraction that is deliberately broader than force-only evaluation.

At publication time, the engine supports several MLIP families—MACE, MatterSim, Fairchem models, Metatensor models, and SevenNet—as well as classical interatomic potentials including Lennard-Jones, Morse, and soft-sphere (Cohen et al., 8 Aug 2025). The coexistence of learned and analytic potentials is significant because it places TorchSim simultaneously in the lineage of differentiable simulation frameworks and in the more conventional tradition of general-purpose atomistic engines. Related differentiable molecular simulation systems, such as TorchMD (Thölke et al., 2022) and JAX MD (Schoenholz et al., 2019), situate TorchSim within a broader movement toward tensor-framework-native simulation, but the TorchSim paper emphasizes batched multi-system execution and materials-informatics interoperability as its distinguishing combination of features (Cohen et al., 8 Aug 2025).

3. Batched multi-system execution

TorchSim treats batching as its defining systems innovation. Instead of executing one structure per forward pass, it evolves many independent systems simultaneously (Cohen et al., 8 Aug 2025). The reported motivation is that even a large GPU such as an H100 can be severely underutilized by single-system execution for some MLIP workloads, whereas batching substantially improves throughput.

The internal representation is based on a single SimState. Multiple systems are represented by concatenating atom-wise quantities such as positions, velocities, and masses, while stacking batch-wise quantities such as unit cell and stress tensor (Cohen et al., 8 Aug 2025). This arrangement allows systems of different sizes to be advanced together in one batched simulation state. The abstraction is not merely a container format: it is the data model that enables mixed-size concurrent evolution.

The paper formulates batching as a memory-constrained packing problem. Systems $S_i$ are packed into bins subject to the constraint

$\sum_{i \in \text{bin}} m_i \le M_{\max}$

where $m_i$ is a model-specific memory scaler and $M_{\max}$ is the empirically determined capacity (Cohen et al., 8 Aug 2025). The source stresses that $m_i$ is not the literal memory footprint but a surrogate determined by model characteristics, for example $n_{\text{atoms}}$ for fixed-max-neighbor models or $n_{\text{atoms}} \times \text{density}$ for radial-cutoff models.

This emphasis on surrogate-based memory scaling reflects a general property of MLIP workloads: memory demand is architecture-dependent and often dominated by graph construction or neighbor-list behavior rather than by the integrator itself. A plausible implication is that TorchSim’s batching machinery is intended to be model-aware rather than merely device-aware.

4. AutoBatching and memory management

TorchSim introduces AutoBatching, a mechanism intended to organize simulations so that they fit within GPU memory while preserving high utilization (Cohen et al., 8 Aug 2025). The system uses a model-specific memory_scales_with parameter to estimate memory in the appropriate unit. The paper notes that some models scale with atom count, while others depend on atom count times density or on graph-neighbor construction details such as maximum neighbor count.

Maximum device capacity is estimated empirically by running progressively larger forward passes until an out-of-memory error occurs; this capacity can be discovered on the fly or cached for reuse (Cohen et al., 8 Aug 2025). The accompanying conceptual workflow is described as: determine memory scalers for systems, select largest and smallest systems, iteratively add systems until OOM, and conservatively compute the maximum supported memory scaler.

Two autobatching modes are defined.

Mode	Intended use	Mechanism
BinningAutobatcher	Fixed-length simulations	1D bin-packing of systems into memory-bounded batches
InFlightAutobatcher	Variable-length simulations	Streaming replacement as systems converge

BinningAutobatcher is used when the simulation length is known in advance, such as running 100 ps of molecular dynamics on 1000 systems (Cohen et al., 8 Aug 2025). It computes each system’s memory footprint, packs systems into bins of size equal to the model’s maximum memory, and runs those bins sequentially. InFlightAutobatcher is used for variable-length procedures such as geometry relaxation, where convergence time is not known a priori. It initializes a batch from an input stream until memory is full, runs several optimization steps, removes converged systems, replaces them with new systems, and repeats until all systems converge (Cohen et al., 8 Aug 2025).

This separation between fixed-length and variable-length batching is consequential. It indicates that TorchSim’s batching policy is linked to workflow semantics rather than being a single static scheduler. In relaxation-heavy workloads, where different structures terminate after different numbers of steps, the in-flight replacement strategy is intended to prevent GPU underutilization due to early convergence of a subset of the batch.

5. Simulation capabilities and differentiability

TorchSim supports molecular dynamics integrators, structural relaxation optimizers, classical and machine-learned interatomic potentials, differentiable simulation, automatic memory management, GPU acceleration, and interoperability with materials informatics tools (Cohen et al., 8 Aug 2025). The publication also indicates future expansion to nudged elastic band methods, phonopy-based force-constant calculations, grand canonical Monte Carlo, metadynamics, multi-GPU and multi-node execution, and torch.compile optimization (Cohen et al., 8 Aug 2025). These are framed as future directions rather than present capabilities.

Differentiability is a major architectural consequence of implementing the simulation loop in PyTorch. Because the primitives are tensor operations tracked by autograd, model parameters can in principle be optimized through full trajectories. The paper highlights the case of differentiating with respect to observables computed after molecular dynamics runs, such as radial distribution functions. It presents this schematically as a loss of the form

$$\mathcal{L} = \mathcal{L}\big(\mathrm{RDF}(\text{trajectory}), \mathrm{RDF}_{\text{exp}}\big)$$

differentiated through the simulation pipeline (Cohen et al., 8 Aug 2025). This is an important distinction from workflows in which simulation is merely downstream of model inference and not part of the computational graph.

TorchSim also emphasizes interoperability with existing materials data structures. It can ingest and export pymatgen.core.Structure, ase.Atoms, and phonopy.PhonopyAtoms through conversions to and from its internal SimState representation (Cohen et al., 8 Aug 2025). This matters because it positions the engine not as an isolated software silo but as a component that can be embedded into established materials science pipelines. The relationship to ASE is particularly salient: ASE remains a central workflow layer in computational materials science, but the TorchSim paper argues that ASE’s conventional execution model is not optimized for batched, GPU-native MLIP simulation (Cohen et al., 8 Aug 2025).

6. Benchmarks, throughput metrics, and comparative positioning

The main benchmark reported for TorchSim uses face-centered cubic Cu systems ranging from 16 to 864 atoms on an NVIDIA H100 (Cohen et al., 8 Aug 2025). Four setups are compared: serial Langevin molecular dynamics in ASE, serial Langevin molecular dynamics in TorchSim, single-step MLIP force evaluations in TorchSim, and batched Langevin molecular dynamics in TorchSim. The protocol includes a 15-step warmup followed by 150 molecular dynamics steps or 150 forward passes (Cohen et al., 8 Aug 2025).

The reported maximum capacities for selected models are as follows.

Model	Maximum capacity
MACE-MPA-0	8,000 atoms
EGIP	10,000 atoms
SevenNet-MF-OMPA	3,200 atoms
MatterSim V1 1M	22,000 atoms
PET-MAD	9,000 atoms

The paper’s principal performance conclusion is that the dominant cost is model forward evaluation rather than integrator overhead (Cohen et al., 8 Aug 2025). It states that serial molecular dynamics in TorchSim and ASE have similar performance, and that looping over forward passes is not much faster than performing full molecular dynamics integration, implying that the speedup does not primarily come from a faster integrator. Instead, the gain comes from batching.

The headline result is approximately an order-of-magnitude speedup over ASE on small systems for all tested models, with substantially larger gains in some cases (Cohen et al., 8 Aug 2025). A specific example is that MACE-MPA-0 can run 500 16-atom systems in parallel with throughput about 100× faster than serial execution (Cohen et al., 8 Aug 2025). The paper further argues that atoms per second is the most informative throughput metric because it reflects both system size and batch efficiency. This choice of metric is consequential because it shifts attention away from traditional molecular-dynamics reporting conventions such as nanoseconds per day, which do not compare heterogeneous batched workloads as directly.

In unbatched runs, throughput scales slightly sublinearly with system size, and MACE becomes the fastest at larger sizes despite not being the fastest for the smallest systems (Cohen et al., 8 Aug 2025). In batched runs, system size matters much less, because many systems are processed simultaneously and differences are dominated by graph construction and model-specific overheads rather than per-system dynamics. This suggests that the TorchSim performance model is governed less by classical timestep integration costs and more by the architecture of the potential evaluator.

7. Package organization, ecosystem context, and limitations

The appendix-level package structure described in the paper presents TorchSim as a modular framework with components including models, optimizers, integrators, autobatching, neighbors, io, state, md, nvt, nve, npt, monte_carlo, properties, correlations, math, and transforms, together with model-specific integrations such as mace, sevennet, mattersim, fairchem, graphpes, and metatomic (Cohen et al., 8 Aug 2025). This organization indicates that TorchSim is intended as a general simulation environment rather than as a narrowly specialized batched molecular-dynamics wrapper.

In relation to prior software, the paper positions TorchSim against ASE, LAMMPS, and HOOMD-blue on the grounds that those frameworks are not centered on PyTorch-native MLIPs or differentiable batched workflows (Cohen et al., 8 Aug 2025). It also places itself alongside differentiable simulation frameworks such as TorchMD (Thölke et al., 2022) and JAX MD (Schoenholz et al., 2019), while claiming broader materials-informatics integration and batched multi-system execution support (Cohen et al., 8 Aug 2025). The comparison table summarized in the source states that TorchSim supports batching, diverse MLIPs, differentiability, pure Python, and GPU dynamics, while multi-GPU support is not yet available (Cohen et al., 8 Aug 2025).

That last point is central to the present scope of the engine. The source explicitly notes future work on multi-GPU and multi-node execution, which indicates that the published version is currently oriented toward single-accelerator high-throughput workloads rather than distributed execution (Cohen et al., 8 Aug 2025). This is not a minor engineering omission but a defining limitation of the initial release. A plausible implication is that TorchSim is best understood as a throughput-oriented GPU-native engine for MLIP-heavy simulation and optimization pipelines, with scalability currently derived from batching across many systems rather than from distributed parallelism across large hardware ensembles.