Three critical regressions in MatterSim 1.2.5: TorchSim NaNs, CUDA fine-tune device mismatch, and incorrect batched three-body indices

[sek1ro-yuzzz]

Environment

• MatterSim: 1.2.5
• Python: 3.12.13
• PyTorch: 2.8.0+cu128
• CUDA: 12.8
• GPU: NVIDIA GeForce RTX 5090
• ASE: 3.28.0
• Models tested: mattersim-v1.0.0-1M, mattersim-v1.0.0-5M

I encountered three independent but serious issues while benchmarking and fine-tuning MatterSim.

---

1. MatterSim TorchSimWrapper returns NaN for large static systems

For the same periodic structures, ASE MatterSimCalculator gives finite single-point energy/forces, but MatterSim TorchSimWrapper returns NaN energy/forces for larger systems.

Observed static validation:

model	atoms	ASE energy finite?	TorchSim energy	TorchSim forces finite?
1M	32	yes	finite, matches ASE	yes
1M	128	yes	finite, matches ASE	yes
1M	1024	yes, -4589.7275 eV	nan	no
1M	2304	yes, -10326.9229 eV	nan	no
5M	32	yes	finite, matches ASE	yes
5M	128	yes	finite, matches ASE	yes
5M	1024	yes, -4600.2485 eV	nan	no
5M	2304	yes, -10350.5645 eV	nan	no

This appears before MD propagation, so it is not just an unstable MD trajectory. The static TorchSim wrapper output is already non-finite.

---

2. Official fine-tune path fails on CUDA with CPU/CUDA device mismatch

Running the official fine-tune path on CUDA can fail with:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The failing operation is in the loss path, e.g.

e_pred = result["total_energy"] / graph_batch.num_atoms

Root cause: batch_to_dict(graph_batch, device=self.device) moves model input tensors to CUDA, but loss_calc() still uses labels and metadata from the original graph_batch, including:

graph_batch.energy
graph_batch.forces
graph_batch.stress
graph_batch.num_atoms

In MatterSim 1.2.5 these remain on CPU in the CUDA fine-tune path.

A minimal local workaround was:

@@ train_one_epoch
             else:
+                graph_batch.to(self.device)
                 input = batch_to_dict(graph_batch, device=self.device)

A robust upstream fix should ensure all tensors used by loss_calc() are on the same device as the predictions.

---

3. Incorrect batched three_body_indices in fine-tune dataloader for batch_size > 1

This is the most serious fine-tune correctness bug I found.

In MatterSim 1.2.5, GraphConverter.convert() still returns plain PyG Data:

return Data(**args)

but the code now relies on M3GNetData.__inc__() to offset three_body_indices by num_bonds during PyG collation:

class M3GNetData(Data):
    def __inc__(self, key, value, *args, **kwargs):
        if key == "three_body_indices":
            return self.num_bonds
        return super().__inc__(key, value, *args, **kwargs)

Since the legacy converter returns Data, this __inc__() is never used. Therefore, for batch_size > 1, three_body_indices from later structures point to the wrong global edge indices.

This causes pretrained validation to become strongly batch-size dependent.

Observed on the same AgGaSe validation set with pretrained 5M and official build_dataloader:

batch size	force MAE
1	0.05132 eV/A
2	0.15521 eV/A
4	0.20043 eV/A
8	0.22671 eV/A
16	0.24567 eV/A
32	0.26489 eV/A

If I wrap the same graph objects as M3GNetData, batch_size=32 returns to:

force MAE = 0.05132 eV/A

Minimal local fix:

# mattersim/datasets/utils/converter.py
-            return Data(**args)
+            return M3GNetData(**args)

After this change, batch_size=1 and batch_size=32 agree again, and fine-tune validation matches the ASE calculator baseline.

This also explains why older MatterSim 1.1.x did not show this problem: older m3gnet.py manually applied an index_bias inside forward():

cumsum = torch.cumsum(num_bonds, dim=0) - num_bonds
index_bias = torch.repeat_interleave(cumsum, num_three_body, dim=0).unsqueeze(-1)
three_body_indices = three_body_indices + index_bias

In 1.2.5 that forward-time correction was removed, but the legacy converter was not updated to return M3GNetData.

Expected behavior

• TorchSim static outputs should stay finite whenever ASE MatterSimCalculator gives finite outputs for the same structure/model.
• CUDA fine-tune should not fail due to CPU/CUDA label tensor mismatch.
• Pretrained validation metrics should not depend strongly on dataloader batch_size; batch_size=1 and batch_size=32 should agree within numerical noise.

Suggested fixes

1. Investigate TorchSim large-system graph construction / neighbor handling causing NaNs.
2. Move graph_batch labels/metadata to the training device before loss_calc().
3. Change legacy GraphConverter.convert() to return M3GNetData(**args), or otherwise restore correct three_body_indices offset handling for batched data.