Batched MD with SevenNet+D3 can produce NaN trajectories while ASE remains finite

[sek1ro-yuzzz]

Summary

I am seeing reproducible NaNs during TorchSim batched MD with SevenNet v0.13.0 + D3(BJ)-PBE. The same structures and MD settings remain finite when run through the ASE SevenNet+D3 path.

The TorchSim call follows the documented batched-MD style:

ts.integrate(system=list_of_atoms, ...)

where list_of_atoms contains 16 independent ASE Atoms objects.

System

• 16 independent amorphous Na16Ta16Cl96 structures
• 128 atoms per structure
• Initial structures look reasonable; initial nearest-neighbor distances are about 2.26 Angstrom
• SevenNet OMNI checkpoint with modal="omat24"
• D3(BJ)-PBE enabled

MD settings

• Ensemble: fixed-cell NVT
• Integrator: Nose-Hoover chain
• Temperature: 300 K
• Timestep: 2 fs
• Thermostat damping: 200 fs
• Warmup: 200 steps
• Measured segment: 1000 steps

Environment

Primary test:

• GPU: NVIDIA RTX 5090
• sevenn 0.13.0
• torch 2.11.0+cu128
• torch-sim-atomistic 0.6.0
• ase 3.28.0

Repeat test:

• GPU: NVIDIA RTX 4090
• Same package versions: sevenn 0.13.0, torch 2.11.0+cu128, torch-sim-atomistic 0.6.0, ase 3.28.0

Observed behavior

On the RTX 5090 machine, for the same batch of 16 structures:

• TorchSim + cueq: MD completes, but some batch elements end with NaN energy/temperature.
• TorchSim + flash: MD completes, but some batch elements end with NaN energy/temperature.
• ASE + SevenNet+D3 + flash: all 16 structures remain finite under the same structures, potential, and MD settings.

On the RTX 4090 machine, repeating TorchSim + cueq also produces NaN energy/temperature in some batch elements. This suggests the issue is not specific to a single RTX 5090/Blackwell machine.

Static single-point comparison

For the same structure, TorchSim SevenNet+D3 and ASE SevenNet+D3 agree closely in static single-point calculations:

• Energy difference: about 0.000679 meV/atom
• Max force-component difference: about 1.62e-5 eV/Angstrom
• Max stress-component difference: about 1.13e-8 eV/Angstrom^3

So the problem appears more likely to be in the TorchSim MD state update / batched integration path than in the static SevenNet+D3 energy/force evaluation itself.

Expected behavior

With the same initial structures, potential, and fixed-cell NVT settings, TorchSim batched MD should keep energy, temperature, and coordinates finite. If one batch element becomes numerically unstable, it would also be helpful to have a diagnostic exception or warning rather than silently returning a final state with partial-batch NaNs.

Actual behavior

TorchSim batched MD completes and returns results, but some batch elements end with NaN energy/temperature. The ASE MD control remains finite under the same conditions.

I can provide the benchmark config, structures, or a smaller reproducer if that would be useful.

[CompRhys]

Can you provide code to reproduce? We need to work out if this is internal to torchsim or an implementation issue in either sevennet upstream or potentially alchemiops if using our wrapper of their batched d3 code.

[sek1ro-yuzzz]

torchsim_sevennet_d3_md_check.ipynb
stru.tar.gz

The code and structure I used is as above. It can be found that when d3 is turned off, the problem of NaN still occurs. In addition, before I use mattersim also encountered similar problems, see:https://github.com/microsoft/mattersim/issues/163。

[curtischong]

I think the issue is that we are not properly converting the input parameters of the integrators to values in proper Metal units. We are properly converting the dt value https://github.com/TorchSim/torch-sim/blob/main/torch_sim/runners.py#L307 but not these per-integrator specific params.

The NaN occurs because when we don't multiply tau by MetalUnits.time, the value of tau is roughly 98.23x smaller (In your notebook, tau remains to be 0.2). So every time nose_hoover calls the update_mass function, which is update_chain_mass_fn, we calculate a much smaller Q than we should:

https://github.com/TorchSim/torch-sim/blob/main/torch_sim/integrators/md.py#L549

Q = (
    kT_batched.unsqueeze(-1)
    * torch.square(chain_state.tau).unsqueeze(-1)
    * torch.ones((n_systems, chain_length), dtype=dtype, device=device)
)

I think this causes the NaNs down the line - since the mass of the thermostat (Q) is smaller, it's easier for the system to be unstable and more readily react to position updates.

One way to fix this is to define the proper conversion for each param for each integrator here https://github.com/curtischong/torch-sim/pull/15/changes#diff-a48c24c29550d9df7680a25f296c3547c578901beb3f6ef34e9417ec1de61c74R248

If this is the bug, fixing this is a bit painful because many simulation results may be invalidated, and we'll probably need to announce it. This feels like the bug bc ase requires units like this to be in ase units https://gitlab.com/bigd4/ase/-/blob/master/ase/md/nvtberendsen.py?ref_type=heads#L44. But I'm slightly skeptical bc this seems like a major bug that has been left unnoticed for a while

[CompRhys]

@curtischong How confident are you in the units diagnosis? if you rescale the units you use so they match ase does this resolve the NaN issue? I want to try understand if the problem you found (which I think is indeed a bug) is the cause of these NaNs

[curtischong]

I need to read more of the other workflows in the repo to be more sure if this is even a real problem. sorry I got a bit caught up with other things and was hoping that someone else more familiar could take a second look. Friday is a day off so I'll for sure look more into it then

[curtischong]

I believe that this is the fix. I've reproduced the bug in #586 and in my fix, it disappears curtischong#15.

We didn't convert the units for integrator params in ts.integrate so that PR will fix that! Thank you for reporting!