PERF: Remove GPU-CPU Synchronization in FIRE by curtischong · Pull Request #609 · TorchSim/torch-sim

curtischong · 2026-09-26T13:17:15Z

Summary

There is a lot of GPU-CPU synchronization in the fire optimizer. This PR attempts to remove them all so everything stays on GPU.

Across 5 runs each against main and this branch, Mace small optimized 100 structures with these timings:

Revision	Median time	Time SD	Median structures/min	Throughput SD
main	12.104 s	0.165 s	495.71	6.79
PR (`276d04d`)	11.883 s	0.062 s	504.93	2.66

1.019× throughput: 1.86% faster, or 1.83% less runtime

For faster models, this change will be even faster.

Here are the specific optimizations we implemented:

We remove .item() in batched_vdot
We remove an if statement which caused a sync from GPU to CPU and replaced this with no-branching GPU operations
We no longer call state.n_systems which causes a GPU to CPU sync when we look at the size of the resulting tensor from the .unique() call.

Replaced x[mask] with torch.where calls. Why? It's because:

x[mask] *= f works by trying to ONLY do the minimal amount of multiplication needed. Unfortunately, this means that pytorch needs to dynamically resolve the shape of the indices of x resolved by mask. dynamically resolving the shape of these indices builds a temp array which sends information to the CPU. This is best explained with an example by GPT:

For example, if mask is [True, False, True, False], PyTorch calls nonzero(mask) to produce the index tensor [0, 2] on the GPU.

The CPU knows the mask has 4 entries, but doesn’t know that 2 are true. PyTorch’s CUDA implementation counts the matches, copies that count back to the CPU, and waits for the copy to finish. It can then set the index tensor’s shape to (2,) and gather x[0] and x[2].

torch.where(mask, x*f, x) on the other hand, just pre-multipled all indices of x by f before selecting by the mask. It doesn't need to wait to dynamically figure out the size of the mask. The extra multiplication is cheaper than the wait.

Here is a script to benchmark the two approaches

import time
import torch

mask = torch.tensor([True, False] * 4, device="cuda")
factor = torch.tensor(0.99, device="cuda")


def indexing(x):
    x[mask] *= factor
    return x


def where(x):
    return torch.where(mask, x * factor, x)


def measure(update):
    x = torch.ones(8, device="cuda")

    # Warm up before starting the timer.
    for _ in range(100):
        x = update(x)
    torch.cuda.synchronize()

    start = time.perf_counter()
    for _ in range(1000):
        x = update(x)
    torch.cuda.synchronize()  # Wait for all timed GPU work to finish.

    return (time.perf_counter() - start) / 1000 * 1e6


print(f"indexing: {measure(indexing):.1f} microseconds/update")
print(f"where:    {measure(where):.1f} microseconds/update")

# output:
# indexing: 66.7 microseconds/update
# where:    12.2 microseconds/update

curtischong · 2026-09-26T13:26:13Z


    # Initialize velocities if NaN
    nan_velocities = state.velocities.isnan().any(dim=1)
-    if nan_velocities.any():


this if statement triggered a sync between gpu and cpu

curtischong · 2026-09-26T14:12:36Z

    """Perform one ASE-style FIRE optimization step."""
    from torch_sim.optimizers import CellFireState

-    n_systems, device, dtype = state.n_systems, state.device, state.dtype


n_systems is slow. since we do a torch.unique().shape[0] call. we need to wait for the unique to run.

maybe we should have a proper fix for this by just manually tracking the number of systems in simstate. but this feels like it can easily lead to bugs

can we not look at the state.cell.numel()/9 to avoid the unique call? seems a little easier to follow than state.dt.numel() and could replace the n_systems property on state level maybe?

I agree fixing it at the property level is the better way to go: store_model_extras and compute_cell_forces also read state.n_systems, so that would remove those syncs too.
We could go back to state.n_systems here instead of dt.numel() once that change lands.
Happy to do that in a separate PR if that works for you both.

I think Rhys's idea might be better than what I proposed (manually tracking the number of systems) since the type definition of a simstate says that a cell always exists:

cell (torch.Tensor): Unit cell vectors with shape (n_systems, 3, 3).

So we should be able to remove the torch.unique() by calling state.cell.shape[0]

I was actually surprised that a cell must always exist bc it might be inconvenient for bio systems. But a hack is to just make a super large cell

Opened #621 for this.

JFYI on molecules: ASE gives a molecule a zero cell, and atoms_to_state carries it over as is, I think.

curtischong · 2026-09-26T14:13:37Z

    """Perform one ASE-style FIRE optimization step."""
    from torch_sim.optimizers import CellFireState

-    n_systems, device, dtype = state.n_systems, state.device, state.dtype


maybe we should have a proper fix for this by just manually tracking the number of systems in simstate. but this feels like it can easily lead to bugs

curtischong · 2026-09-26T15:12:10Z

        )


-def test_fire_ase_negative_power_branch(


the parameter-reset and velocity-update checks in this test are now covered by test_fire_mixed_power_update_order

curtischong · 2026-09-26T17:51:30Z

-        if isinstance(state, CellFireState):
-            system_power += (state.cell_forces * state.cell_velocities).sum(dim=(1, 2))
+    # Only update dt, alpha, and n_pos after the first step.
+    pos_mask_system &= ~is_first_step


this is how we get rid of the if statement. if it is the first step, pos_mask_system and neg_mask_system will be 0. so all the downstream variables that only trigger on the else part of the if statement will not be triggered

curtischong · 2026-09-26T17:55:35Z

    if isinstance(state, CellFireState):
        state.cell_velocities.nan_to_num_(nan=0.0)

-    if nan_velocities.all():


this if statement is slow since we need to move the result of this GPU tensor to the CPU

curtischong · 2026-09-26T18:58:23Z

-            (1.0 - alpha_atom) * state.velocities + alpha_atom * v_mixing_atom,
-            torch.zeros_like(state.velocities),
+        state.cell_velocities = torch.where(
+            ~is_first_step, mixed_cell_velocities, state.cell_velocities


only apply the mixed_cell_velocities if we are NOT on the first step (this is us removing the if statement and having non-branching code)

curtischong · 2026-09-26T19:00:31Z

    # Evolve 10 steps so system 0 has non-trivial FIRE state (dt, alpha, n_pos)
    for _ in range(10):
-        state = ts.fire_step(state=state, model=lj_model)
+        state = ts.fire_step(state=state, model=lj_model, fire_flavor=fire_flavor)


adding this increases test coverage for this PR so the codecov check doesn't fail us. previously, ts.fire_step only ran the ase_fire fire_flavor

curtischong · 2026-09-26T20:25:23Z

+    forces = state.forces
+    if isinstance(state, CellFireState):
+        forces = torch.where(is_first_step, forces, state.deform_grad_forces())


this is a bit weird to read but these three lines are the main part of the if/else statements we replaced in the previous code (ecluding the alpha_start_system declaration).

by default, the forces is just state.forces (seen in the if nan_velocities.all(): branch of the if statement, and in the else branch of the if isinstance(state, CellFireState) else state.forces if the first if statement failed.

we need to do the torch.where(is_first_step function which is like "the else branch triggering" in a non-branching way.

curtischong · 2026-09-26T20:30:36Z

-    state.alpha[neg_mask_system] = alpha_start_system[neg_mask_system]
-    state.n_pos[neg_mask_system] = 0
+    state.dt = torch.where(neg_mask_system, state.dt * f_dec, state.dt)
+    state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha)


we got rid of alpha_start_system in this diff becuase we build the state.alpha tensor with this line:

state.alpha = torch.where(inc_mask, state.alpha * f_alpha, state.alpha)

and this line:

state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha)
notice that this second line defaults to state.alpha when neg_mask_system is false

so we are properly replicating all the state.alpha[inc_mask] *= f_alpha and state.alpha[neg_mask_system] = alpha_start_system[neg_mask_system] in the previous code

curtischong · 2026-09-27T13:35:57Z


-        # Calculate power (newly zeroed systems will have power=0 → neg_mask)
-        system_power = tsm.batched_vdot(forces, state.velocities, state.system_idx)
-        if isinstance(state, CellFireState):


unfortunately the github diff from these lines onwards are hard to read. but they follow the original implementation quite closely

lil-lon

LGTM. I reviewed the changes, and they preserve the existing logic (apart from the intended vv_fire cell-velocity NaN fix) while removing CPU/GPU syncs, including:

FIRE first-step detection
batched_vdot
updates of dt, alpha and n_pos

The way the number of systems is obtained can be improved as discussed in the thread above, but that can be done in a follow-up.
If @CompRhys is okay with this change, I think this PR is ready to merge.

CompRhys · 2026-10-21T12:03:27Z

If @CompRhys is okay with this change, I think this PR is ready to merge.

Feel free to merge!

lil-lon

LGTM!

lil-lon · 2026-10-22T00:57:54Z

@CompRhys Thanks for checking! I don't have permission to merge into main, so could you merge this one?

curtischong added 2 commits September 26, 2026 06:17

v1 pass

cb17b38

rm slop

2e414a1

curtischong commented Sep 26, 2026

View reviewed changes

curtischong changed the title ~~v1 pass~~ Remove GPU-CPU Synchronization in FIRE Sep 26, 2026

curtischong commented Sep 26, 2026

View reviewed changes

curtischong added 10 commits September 26, 2026 07:41

simpify hopefully

9d698f0

smaller diff

72c5f1e

simplify tests

6fc4eb0

simplify tests more

3bda54e

more cleanup

5e8e03a

rm tests

25982c2

maybe cleaner

b58c15a

maybe cleaner

5df530b

revert to clear form

a58ea54

maybe

1d11b9d

curtischong commented Sep 26, 2026

View reviewed changes

more test coverage

8c7d11c

curtischong commented Sep 26, 2026

View reviewed changes

curtischong changed the title ~~Remove GPU-CPU Synchronization in FIRE~~ PERF: Remove GPU-CPU Synchronization in FIRE Sep 27, 2026

curtischong commented Sep 27, 2026

View reviewed changes

curtischong added 2 commits September 27, 2026 06:47

fixes

f2c13e1

simplify test

276d04d

curtischong marked this pull request as ready for review September 27, 2026 15:29

lil-lon reviewed Oct 21, 2026

View reviewed changes

lil-lon approved these changes Oct 21, 2026

View reviewed changes

Merge branch 'main' into rm-cpu-transfers

2027090

CompRhys enabled auto-merge (squash) October 22, 2026 10:07

CompRhys merged commit 3cbb5dc into TorchSim:main Oct 22, 2026
64 checks passed

curtischong deleted the rm-cpu-transfers branch October 22, 2026 11:27

lil-lon mentioned this pull request Oct 24, 2026

perf(state): derive n_systems from cell shape to avoid GPU-CPU sync #621

Merged

3 tasks done

Conversation

curtischong commented Sep 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong Sep 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lil-lon left a comment

Choose a reason for hiding this comment

Uh oh!

CompRhys commented Oct 21, 2026

Uh oh!

lil-lon left a comment

Choose a reason for hiding this comment

Uh oh!

lil-lon commented Oct 22, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

curtischong commented Sep 26, 2026 •

edited

Loading

curtischong Oct 22, 2026 •

edited

Loading

curtischong Sep 26, 2026 •

edited

Loading