SevenNet+D3 with TorchSim batched MD can produce NaNs while ASE remains finite

[sek1ro-yuzzz]

Summary

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

I also opened a TorchSim-side issue here:

TorchSim/torch-sim#579

I am reporting this here as well in case the SevenNet TorchSim wrapper, accelerator configuration, or SevenNet+D3 combination has known caveats or recommended diagnostics.

Setup

The TorchSim call follows the official 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
• modal="omat24"
• D3(BJ)-PBE enabled

MD settings:

• Fixed-cell NVT
• Nose-Hoover chain
• 300 K
• 2 fs timestep
• 200 fs thermostat damping
• 200 warmup steps
• 1000 measured MD 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 + SevenNet+D3 + cueq: MD completes, but some batch elements end with NaN energy/temperature.
• TorchSim + SevenNet+D3 + 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 + SevenNet+D3 + 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 MD / batched integration path than in the static SevenNet+D3 energy/force evaluation itself.

Question

Do you have any suggestions for diagnosing this from the SevenNet side? In particular:

• Are there known limitations for SevenNet v0.13.0 with TorchSim batched MD and D3?
• Are cueq / flash expected to be numerically equivalent for MD use here?
• Is there a recommended minimal diagnostic to separate a SevenNet wrapper issue from a TorchSim integration/state-update issue?

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

[alphalm4]

Thanks for reporting this. before working further, what is the batch-D3 implmenetation you used? nvalchemiops (by torchsim) or our own SevenNetD3Model?

[sek1ro-yuzzz]

Thanks for checking. I used SevenNetD3Model, i.e. your own implementation, not nvalchemiops.

(Also, thank you for the D3 support in SevenNet. I found it very powerful, and among the D3 modules I have tried, it has been the fastest so far.)

That said, my current guess is that the issue may be related to torchsim rather than the D3 implementation itself, because I have encountered similar problems even when D3 is disabled, and also when using MatterSim.

[YutackPark]

Are there known limitations for SevenNet v0.13.0 with TorchSim batched MD and D3?

No, the batched D3 in SevenNet was first introduced in v0.13.0. So, this is the only possible version of 7net that can test this setup.

Are cueq / flash expected to be numerically equivalent for MD use here?

Numerical equivalence is usually not guaranteed when a GPU is involved (in any set of settings). However, both kernels can be utilized for batched evaluation without a problem (for instance, both can be used for training).

Is there a recommended minimal diagnostic to separate a SevenNet wrapper issue from a TorchSim integration/state-update issue?

Below is where SevenNet and torchSim connected. SevenNet model batched inference entirely happens in SevenNet side and returns energy, force, and stress as promised by the torchsim. No other hidden interactions, which alter results, as far as I know:

TorchSim side: https://github.com/TorchSim/torch-sim/blob/87dad4d340a372f70d0e30eb06b0f28aa4d473b7/torch_sim/models/interface.py#L98-L101

SevenNet side:

SevenNet/sevenn/torchsim.py

Lines 479 to 500 in 3a7552c

	results['energy'] += torch.from_numpy(d3_energy).to(
	device=self._device, dtype=self._dtype,
	)
	results['forces'] += torch.from_numpy(
	np.ascontiguousarray(d3_forces),
	).to(device=self._device, dtype=self._dtype)
	# D3 stress from batch kernel is extensive virial [B, 3, 3] in eV
	# ASE/TorchSim convention: stress = -virial / volume (eV/A^3)
	# For non-periodic systems (zero cell -> zero volume)
	# stress is physically undefined, so we zero it out.
	volumes = torch.det(state.row_vector_cell.detach()).abs().cpu().numpy()
	has_volume = volumes > 0
	safe_volumes = np.where(has_volume, volumes, 1.0)
	d3_stress_intensive = np.where(
	has_volume[:, None, None],
	-d3_stress / safe_volumes[:, None, None],
	0.0,
	)
	results['stress'] += torch.from_numpy(
	np.ascontiguousarray(d3_stress_intensive),
	).to(device=self._device, dtype=self._dtype)

---

P.S. Do you get NaNs from the first MD step? If not, are they (batched MD and ASE calculator MD) returning the same energy, forces before it returns NaN?

[sek1ro-yuzzz]

We already checked static parity between ASE SevenNet+D3 and TorchSim SevenNet+D3 on the same Packmol NaTaCl6 structure. The outputs agree closely: energy difference is about 0.000679 meV/atom, max force-component difference is about 1.62e-5 eV/Angstrom, and max stress-component difference is about 1.13e-8 eV/Angstrom^3.

So the NaNs do not appear to come from the SevenNet static batched inference itself. They appear during MD integration. This is also consistent with the TorchSim-side discussion, where a missing MetalUnits.time conversion for Nose-Hoover tau was identified as a possible cause.

[YutackPark]

@sek1ro-yuzzz
Thanks for the information. If I understand correctly, this is not an issue on the 7net side. I'll keep this issue open for other users. Let us know if there's anything on 7net's side worth checking.