A1 Batched relaxation for screening¶

Objectives

Explain why batched GPU relaxation speeds up screening.
Relax one set serially with ASE and in one TorchSim batch.
Compare wall time and energy agreement.

Serial relaxation leaves the GPU idle; batching fills it (GPU engines).

One MACE-MP model, two routes: serial ASE and one batched TorchSim call on one GPU.
Part A: own pixi environment, MACE-MP-0b small (not Part B’s MACE-MP-0a).
Workload: rattled copies of Cu fcc, Si diamond, Fe bcc and Al fcc (8 to 32 atoms); each copy needs its own relaxation.

Serial ASE relaxation versus one batched TorchSim call.

Serial and batched runs¶

Both routes use FIRE with a fixed cell.

Stop at --fmax (default 0.05 eV/Å) or --max-steps.
Serial relaxes only the first --baseline-n structures and extrapolates to the full set; a full serial run would use most of the allocation.

TorchSim advances all ASE Atoms as one batched state; converged structures leave the batch.
--autobatch splits the set to fit GPU memory.
Same checkpoint for both; TorchSim needs the raw model (return_raw_model=True):

models.py

def load_mace(name: str, device: torch.device, dtype: torch.dtype,
              checkpoint: str | None = None) -> ModelPair:
    from mace.calculators.foundations_models import mace_mp
    from torch_sim.models.mace import MaceModel, MaceUrls

    source = checkpoint or getattr(MaceUrls, MACE_MODELS[name]).value
    options = {"model": source, "device": str(device),
               "default_dtype": str(dtype).removeprefix("torch.")}

    def make_torchsim():
        raw = mace_mp(return_raw_model=True, **options)
        return MaceModel(model=raw, device=device, dtype=dtype,
                         compute_forces=True, compute_stress=False)
    return ModelPair(keeps_dtype(make_torchsim),
                     keeps_dtype(lambda: mace_mp(**options)), checkpoint)


def orb_ase_calc(orbff, adapter, device: torch.device, dtype: torch.dtype):
    """ORBCalculator builds its input graph in the global default dtype."""
    from orb_models.forcefield.inference.calculator import ORBCalculator

    class Calculator(ORBCalculator):
        def calculate(self, *args, **kwargs):
            with default_dtype(dtype):
                super().calculate(*args, **kwargs)
    return Calculator(orbff, adapter, device=device)


def load_orb(name: str, device: torch.device, dtype: torch.dtype,
             checkpoint: str | None = None, d3: bool = False) -> ModelPair:
    from orb_models.forcefield import pretrained
    from orb_models.forcefield.inference.d3_model import AlchemiDFTD3, D3SumModel
    from torch_sim.models.orb import OrbModel

    loader = getattr(pretrained, ORB_MODELS[name])
    options = {"device": device, "precision": ORB_PRECISION[dtype], "compile": False}
    if checkpoint:
        options["weights_path"] = checkpoint
    with default_dtype():
        orbff, adapter = loader(**options)
    if d3:
        orbff = D3SumModel(orbff, AlchemiDFTD3(functional="PBE", damping="BJ"))
    return ModelPair(lambda: OrbModel(orbff, adapter, device=device, dtype=dtype),
                     lambda: orb_ase_calc(orbff, adapter, device, dtype), checkpoint)


def load_matgl(name: str, device: torch.device, dtype: torch.dtype) -> ModelPair:
    """ASE only: TorchSim has no MatGL model, and MatGL runs in float32."""
    import matgl
    from matgl.ext.ase import PESCalculator

    if dtype != torch.float32:
        raise ValueError("MatGL models run in float32; pass --dtype float32")
    with default_dtype():
        potential = matgl.load_model(MATGL_MODELS[name]).to(device)
    return ModelPair(None, lambda: PESCalculator(potential, stress_unit="eV/A3"))

Outputs: ASE versus TorchSim energies on the shared structures, summary.json and relaxed.extxyz. A local --checkpoint file also gets its SHA-256 recorded.

Laptop run¶

Package examples/torchsim/ with its own pixi environment; entry point examples/torchsim/__main__.py.
smoke uses a toy Lennard-Jones potential: no download, seconds to run, no physical meaning.

cd examples/torchsim
pixi run smoke

Output goes to examples/torchsim_smoke/. Existing output directories are never overwritten: remove it or pass a new --outdir.

Expected output (numbers vary):

model=lj device=cpu dtype=float64 structures=8 serial subset=2
                             wall time   s/structure
serial, measured                 0.1 s        0.03 s   (2 structures)
serial, estimated                0.2 s        0.03 s   (8 structures)
batched, measured                5.4 s        0.68 s   (8 structures)
estimated serial / batched: 0.04
max |E_ASE - E_TorchSim| on the serial subset: 2.03e-01 eV

Batching is slower on a CPU with a cheap potential: no GPU parallelism.
The energy gap comes from different Lennard-Jones cutoff and shift conventions in ASE and TorchSim. Only the MACE run is a consistency check (a few meV).

MACE on a CPU: the cpu task relaxes four structures for at most 20 steps into examples/torchsim_cpu/. First run downloads MACE-MP-0b small (about 68 MB):

pixi run cpu

Leonardo run (one A100)¶

Compute nodes have no internet. On a login node, install the pixi environment in your lesson copy and download the checkpoint to a private location outside Git:

The job checks this SHA-256 and refuses any other file. Submit from the repository root; missing inputs or an existing output directory abort it:

This script has not yet been qualified. Neither result row below comes from it: float64 is from an earlier script with the same workload and options; float32 is from a separate tuned run (options under the table).

Leonardo results¶

One A100, MACE-MP-0b small, no cuEquivariance. Serial measured 8 structures, scaled to the full set.

Precision and batching	Structures	Serial, 8 measured (s)	Serial, estimated (s)	Batched (s)	Estimated serial / batched
float64	64	37.4	299.5	59.3	5.05
float32, `--autobatch`	128	38.7	619.6	85.2	7.27

The float32 run used --dtype float32 --n-variants 32 --autobatch.

Warning

Single examples, not a benchmark.

One run per row: no repeats, medians or ranges.
Full-set serial time is extrapolated from 8 structures.
MACE-MP-0b small, not Part B’s MACE-MP-0a: energies and costs are not interchangeable.
No cuEquivariance kernels; they would change the timings.
Relaxation, not MD: do not compare with Part B timings.

LUMI (AMD MI250X) check¶

MACE runs on AMD via ROCm PyTorch (CSC pytorch module). Separate ASE script, no TorchSim batching, one MI250X GCD:

Step	System	Time (s)
Single point	32 atoms	1.6
FIRE relaxation, serial	6 structures (Cu, Al, Fe, Si)	17.0
MD, 600 K	200 steps, 32 atoms	3.9

One run, float64, MACE-MP-0a small.
Shows the model runs on AMD; not a speed comparison with the A100 table.
CUDA-only kernels (cuEquivariance, ALCHEMI) are not available on AMD.

Reading the results¶

One A100: batched about 5 times faster than serial estimate (float64), about 7 times (float32, --autobatch).
CPU with the toy potential: batching slower; the gain is GPU parallelism.
One run each, extrapolated serial time: ratios are indicative.

Keypoints

Batching runs many relaxations in one GPU call; 5 to 7 times faster on an A100.
On a CPU, batching gives no gain.
Check serial and batched energies agree before comparing timings.