Possible VRAM Control Issue

[alan-he-494165]

I was running the following code to optimise 500 Zn crystal structures

state = ts.initialize_state(compound_init, device=device, dtype=torch.float64)
state_list = state.split()
memory_metric_values = [
    calculate_memory_scaler(s, memory_scales_with="n_atoms_x_density") for s in state_list
]
max_memory_metric = estimate_max_memory_scaler(
    mace_model, state_list, metric_values=memory_metric_values
)
print("Max memory metric", max_memory_metric)

batcher = InFlightAutoBatcher(
    mace_model,
    max_memory_padding=1,
    max_memory_scaler=max_memory_metric * 0.8
)

convergence_fn = ts.generate_force_convergence_fn(0.025)
relaxed_state = ts.optimize(
    system=state,
    model=mace_model,
    optimizer=ts.frechet_cell_fire,
    autobatcher=batcher,
    max_steps=1000,
    convergence_fn=convergence_fn,
)

however error below was observed

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[9], line 22
     15 batcher = InFlightAutoBatcher(
     16     mace_model,
     17     max_memory_padding=1,
     18     max_memory_scaler=max_memory_metric * 0.8
     19 )
     21 convergence_fn = ts.generate_force_convergence_fn(0.025)
---> 22 relaxed_state = ts.optimize(
     23     system=state,
     24     model=mace_model,
     25     optimizer=ts.frechet_cell_fire,
     26     autobatcher=batcher,
     27     max_steps=1000,
     28     convergence_fn=convergence_fn,
     29 )
     30 # extract the final energy from the trajectory file
     31 print(relaxed_state.energy)

File [~/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/runners.py:439](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/runners.py#line=438), in optimize(system, model, optimizer, convergence_fn, trajectory_reporter, autobatcher, max_steps, steps_between_swaps, pbar, **optimizer_kwargs)
    436     pbar_kwargs.setdefault("disable", None)
    437     tqdm_pbar = tqdm(total=state.n_batches, **pbar_kwargs)
--> 439 while (result := autobatcher.next_batch(state, convergence_tensor))[0] is not None:
    440     state, converged_states, batch_indices = result
    441     all_converged_states.extend(converged_states)

File [~/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/autobatching.py:1062](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/autobatching.py#line=1061), in InFlightAutoBatcher.next_batch(self, updated_state, convergence_tensor)
   1060 if updated_state.n_batches > 0:
   1061     next_states = [updated_state, *next_states]
-> 1062 next_batch = concatenate_states(next_states)
   1064 if self.return_indices:
   1065     return next_batch, completed_states, self.current_idx

File [~/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/state.py:839](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/X/lib/python3.11/site-packages/torch_sim/state.py#line=838), in concatenate_states(states, device)
    836 # Concatenate collected tensors
    837 for prop, tensors in per_atom_tensors.items():
    838     # if tensors:
--> 839     concatenated[prop] = torch.cat(tensors, dim=0)
    841 for prop, tensors in per_batch_tensors.items():
    842     # if tensors:
    843     concatenated[prop] = torch.cat(tensors, dim=0)

TypeError: expected Tensor as element 1 in argument 0, but got NoneType

---

This sort of type error may be common when cuda memory was heavily used, and nvidia-smi supported this may be the reasoning.

As I have tried, when I added VRAM cleaning commands when manually written all the batching and optimisation cycles as introduced in the tutorial and using subprcess for optimisation when needed - understanding the autobatching, this problem would resolve at this calculation scale. Also, the convergence becomes easier for smaller scale calculation, interestingly.

[alan-he-494165]

At the same time, I am not 100 % sure if the code was correct. Any constructive suggestions would be appreciated!

[alan-he-494165]

PS. It worked well when optimising 50 structures in one go

[CompRhys]

I believe this is related to the issues #219. We are aware of it and working out what we want to do to resolve.

[curtischong]

I believe that the #219 PR works. If you are able to point your repo to that branch you should be able to simulate larger systems.

We are actively trying to improve the repo! After we get the entire repo properly typed these kinds of issues will definitely be caught sooner.

[alan-he-494165]

Thanks for the reply and constant effort on maintaining this repo! #219 would explain the convergence difficulty. I think it is also relevant to #129 , as if I scale it up even further, despite protection mechanism

max_memory_scaler=max_memory_metric * 0.8

was implemented, it would still result in OOM error instead of type error.

[alan-he-494165]

Thank you for pointing towards PR #219, after implementing that the convergence was resolved. However, it still crashes when I was doing large-scale optimsation

---

IndexError                                Traceback (most recent call last)
Cell In[5], line 21
     14 batcher = InFlightAutoBatcher(
     15     mace_model,
     16     max_memory_padding=1,
     17     max_memory_scaler=max_memory_metric * 0.8
     18 )
     20 convergence_fn = ts.generate_force_convergence_fn(0.01)
---> 21 relaxed_state = ts.optimize(
     22     system=state,
     23     model=mace_model,
     24     optimizer=ts.frechet_cell_fire,
     25     autobatcher=batcher,
     26     max_steps=500,
     27     convergence_fn=convergence_fn,
     28 )

File [~/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/runners.py:452](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/runners.py#line=451), in optimize(system, model, optimizer, convergence_fn, trajectory_reporter, autobatcher, max_steps, steps_between_swaps, pbar, **optimizer_kwargs)
    449 for _step in range(steps_between_swaps):
    450     last_energy = state.energy
--> 452     state = update_fn(state)
    454     if trajectory_reporter:
    455         trajectory_reporter.report(state, step, model=model)

File [~/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/optimizers.py:1619](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/optimizers.py#line=1618), in _ase_fire_step(state, model, dt_max, n_min, f_inc, f_dec, alpha_start, f_alpha, max_step, eps, is_cell_optimization, is_frechet)
   1616     state.positions = state.positions + dr_atom
   1618 # 7. Force [/](http://localhost:8895/) stress refresh & new cell forces
-> 1619 results = model(state)
   1620 state.forces = results["forces"]
   1621 state.energy = results["energy"]

File [~/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch/nn/modules/module.py:1739](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch/nn/modules/module.py#line=1738), in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File [~/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch/nn/modules/module.py:1750](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch/nn/modules/module.py#line=1749), in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File [~/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/models/mace.py:312](http://localhost:8895/lab/workspaces/auto-P/tree/Alan_H/Torch_sim_test/Notebooks/miniconda3/envs/alanhe/lib/python3.11/site-packages/torch_sim/models/mace.py#line=311), in MaceModel.forward(self, state)
    308 batch_mask = state.system_idx == b
    309 # Calculate neighbor list for this system
    310 edge_idx, shifts_idx = self.neighbor_list_fn(
    311     positions=state.positions[batch_mask],
--> 312     cell=state.row_vector_cell[b],
    313     pbc=state.pbc,
    314     cutoff=self.r_max,
    315 )
    317 # Adjust indices for the batch
    318 edge_idx = edge_idx + offset

IndexError: index 109 is out of bounds for dimension 0 with size 109

---

Any constructive suggestions would be appreciated!

[curtischong]

That is annoying. I'll run your example code after work. This looks like an autobatcher issue.

[alan-he-494165]

That is annoying. I'll run your example code after work. This looks like an autobatcher issue.

Thank you for your support. Just to note, 500 starting random structure of Zn (with no more than 4 atoms per unit cell) was directly read into the variable: compound_init

[curtischong]

I ran out of time today recreating this (curtischong#3). but I'll get to it tomorrow!

[curtischong]

I'm having trouble replicating your error (although the minimum system size I'm using is 2 atoms - not 1). Maybe try running my branch and seeing if it errors on your machine? curtischong#3

First install the dependencies
uv sync --extra io --extra test --extra mace

Then run
uv run examples/scripts/500_err.py

[alan-he-494165]

Sorry for the delayed reply, I was on a personal leave

I have been trying to running your branch but when installing the dependencies I do found out the following error

error: Distribution `torch==2.7.1 @ registry+https://pypi.org/simple` can't be installed because it doesn't have a source distribution or wheel for the current platform

hint: You're on Linux (`manylinux_2_17_x86_64`), but `torch` (v2.7.1) only has wheels for the following platforms: `manylinux_2_28_aarch64`, `manylinux_2_28_x86_64`, `macosx_11_0_arm64`, `win_amd64`; consider adding your platform to `tool.uv.required-environments` to ensure uv resolves to a version with compatible wheels

Do you think this is fatal? I think in the .toml file it declared torch version >= 2 shall be fine, whereas we got 2.6.0

[CompRhys]

Sorry for the delayed reply, I was on a personal leave

I have been trying to running your branch but when installing the dependencies I do found out the following error

error: Distribution `torch==2.7.1 @ registry+https://pypi.org/simple` can't be installed because it doesn't have a source distribution or wheel for the current platform

hint: You're on Linux (`manylinux_2_17_x86_64`), but `torch` (v2.7.1) only has wheels for the following platforms: `manylinux_2_28_aarch64`, `manylinux_2_28_x86_64`, `macosx_11_0_arm64`, `win_amd64`; consider adding your platform to `tool.uv.required-environments` to ensure uv resolves to a version with compatible wheels

Do you think this is fatal? I think in the .toml file it declared torch version >= 2 shall be fine, whereas we got 2.6.0

This is entirely unrelated to torch-sim so i'm not sure we can help you with it. Lots of the things we depend on will have chosen to deprecate support for the 2.17 glibc compiler and we're only concerned with supporting things that are widely supported. I would work out whether or not you can update your linux version to a more modern build but we're not able to help you with that.

[curtischong]

We've just merged a PR that makes the scope of each SimState attribute to be defined. so my branch can now be ran with systems of 1 atom. Hopefully we'll catch it then