diff --git a/README.md b/README.md
index 917fa21..ec2a73e 100644
--- a/README.md
+++ b/README.md
@@ -101,7 +101,7 @@ three fields and silently ignores everything else (e.g. `Z:I:1`):
| `max_grad_norm` | `10.0` | Gradient clipping threshold |
| `lr_scheduler` | `plateau` | LR schedule — `plateau` (ReduceLROnPlateau) or `step` (StepLR). Stage 1 and Stage 2 share this mode |
| `scheduler_patience` | `15` | For `plateau`: epochs without improvement before LR reduction. For `step`: epoch interval between LR reductions |
-| `early_stop` | `0` | Stop early if the monitored loss does not improve for this many epochs (`0` = off). The monitored loss is the **validation** loss when `valid_file`/`valid_ratio` is set, otherwise the training loss. Per-stage (MACE-style): a stage-1 plateau with `stage2 1` configured jumps straight into Stage 2 instead of ending the run — only a plateau in the final stage terminates it (the advanced stage-2 start is kept across resume). Set it **larger than `scheduler_patience`** so the LR gets a chance to decay first |
+| `early_stop` | `0` | Stop if the monitored loss (validation loss when a validation set is used, else training loss) hasn't improved for N epochs (`0` = off). Per-stage: a stage-1 plateau jumps into Stage 2 instead of ending the run. Use a value larger than `scheduler_patience` |
| `scheduler_factor` | `0.7` | LR reduction factor — multiplied on each decay in both modes |
| `stage2` | `0` | Enable Stage 2 (`1` = on) |
| `start_stage2` | 50 % of epochs | Epoch to switch to Stage 2 |
@@ -121,10 +121,10 @@ function (`train_nep` / `train_nep_sharded`):
|---|---|---|
| `device` | auto | `"cuda"` / `"xpu"` / `"mps"` / `"cpu"`; any other stream-based PyTorch accelerator should also work if passed explicitly |
| `precision` | `"float32"` | dtype for training + store, `"float32"` or `"float64"` |
-| `backend` | `"auto"` | `"loop"`, `"bmm"`, or `"auto"` |
+| `backend` | `"auto"` | `"loop"`, `"bmm"`, or `"auto"`. Auto resolves to `bmm` under `use_compile` (fuses best) and to `loop` in eager mode unless there are ≥20 element types — benchmarks show eager `loop` wins clearly up to ~16 types |
| `use_autograd_forces` | `False` | autograd-through-rij |
| `use_swa` | `False` | maintain SWA-averaged model and save `nep_average.txt` |
-| `use_compile` | `False` | `torch.compile` the analytical compute (faster epochs after a one-time compile; needs Triton; ignored on the autograd path) |
+| `use_compile` | `False` | `torch.compile` the compute (faster epochs after a one-time compile; needs Triton). |
| `print_interval` | `10` | log to screen every N epochs |
| `checkpoint_interval` | `100` | save `checkpoint.pt` every N epochs |
| `prediction_interval` | `20` | every N epochs run predict with the current-epoch weights and overwrite `{energy,force,virial}_train.out` |
@@ -138,7 +138,6 @@ function (`train_nep` / `train_nep_sharded`):
| `run_seed` | `None` | master RNG seed. `None` = random each run; an int makes the run reproducible (weight init + batch shuffle). Saved in `checkpoint.pt`, restored on resume |
| `valid_file` | `None` | validation `.xyz`, `nep_best` and the plateau LR schedule follow the validation loss; writes GPUMD-style `*_test.out` |
| `valid_ratio` | `None` | hold out this fraction (e.g. `0.1`) of `data_file` as the validation set; the split is drawn from `run_seed` and preserved on resume. Mutually exclusive with `valid_file` |
-| `stream_mode` | `False` | keep the dataset in host memory and stream only the current batch to the GPU (basis computed on the fly, CPU batch assembly prefetched one batch ahead). GPU memory scales with `batch` instead of dataset size — use for datasets that don't fit on the card. Numerically identical to the default; costs a modest per-epoch slowdown in eager mode (speed parity under `use_compile`). Works in both `train_nep` and `train_nep_sharded` (each rank streams its own shard) |
---
@@ -367,8 +366,9 @@ The `torchnep/` package is organised as follows:
| `ops.py` | Core differentiable kernels — Chebyshev/angular basis, descriptors, ANN evaluation, ZBL; pure-PyTorch `loop`/`bmm` backends |
| `nep.py` | `NEPCalculator` — loads a `nep.txt` and computes energy/forces/virial/descriptors for single structures |
| `predict.py` | Batched full-dataset inference (`predict_dataset`), writing GPUMD-compatible `*_train.out` files |
-| `train.py` | Single-GPU/CPU training (`train_nep`): data store, two-stage loop, schedulers, checkpoint/restart, periodic predict |
+| `train.py` | Single-GPU/CPU training (`train_nep`): host-resident streaming data store (`StreamDataStore` + prefetching `iter_collated`), two-stage loop, schedulers, checkpoint/restart, periodic predict |
| `train_sharded.py` | Data-sharded multi-GPU/multi-node training (`train_nep_sharded`) via DDP |
+| `compiled_autograd.py` | `torch.compile` for the autograd force path: the first-order dE/drij gradient is materialized into the graph with `make_fx`, so `use_autograd_forces=True` + `use_compile=True` runs one fused dynamic-shape graph instead of an uncompilable double backward |
| `ase_calculator.py` | ASE `Calculator` wrapper (`NEP`) for relaxation, MD, EOS, phonons, … |
| `constants.py` | Shared constants — element table, covalent radii, NEP polynomial coefficients |
diff --git a/releaseNotes.md b/releaseNotes.md
index ff5c120..ff81b2d 100644
--- a/releaseNotes.md
+++ b/releaseNotes.md
@@ -2,19 +2,17 @@
## Unreleased
-- **Per-stage `early_stop`** (MACE-style): a stage-1 plateau with `stage2 1`
- configured now jumps straight into Stage 2 at the next epoch instead of
- terminating the run; only a plateau in the final stage stops training. The
- advanced stage-2 start epoch is saved in the checkpoint, so resumed runs
- stay in Stage 2.
-
-- **`stream_mode`** (`train_nep` and `train_nep_sharded`): keep the dataset
- (or each rank's shard) in host memory and stream only the current batch to
- the GPU, computing the Chebyshev / angular basis on the fly (CPU batch
- assembly prefetched one batch ahead). GPU memory scales with `batch`
- instead of dataset size. Eager runs are bit-identical to the default
- preloaded mode; under `use_compile=True` the per-batch basis is compiled
- too.
+- **Streaming-only data path**: the preloaded GPU data store and the
+ `stream_mode` option are removed — the dataset stays in host memory and
+ batches are streamed to the device. Same speed, ~10–15x less GPU memory.
+- **`backend="auto"`**: eager mode now uses `loop` below 20 element types
+ (was 8); `bmm` under `use_compile`.
+- **Compiled autograd forces**: `use_autograd_forces=True` +
+ `use_compile=True` now works (first-order gradient materialized via
+ `make_fx`) — ~4x faster than eager autograd.
+- **Per-stage `early_stop`**: a stage-1 plateau jumps into Stage 2 instead
+ of ending the run; only a final-stage plateau stops training (kept across
+ resume).
## 1.0.1
diff --git a/tests/README.md b/tests/README.md
index 76a28ab..dbdb5e9 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -16,7 +16,9 @@ TEST_DEVICE=cpu pytest tests/ # restrict device (default: cpu + cuda if pre
| `test_parsing.py` | Legacy and current `l_max` nep.in / nep.txt parsing. |
| `test_ase_calculator.py` | Optional ASE calculator (energy/forces/stress, ZBL split). |
| `test_b1_and_gpumd_qscaler.py` | Analytical `b1` offset (residual → 0, `nep_best` ≤ `nep_final`); `use_gpumd_qscaler` reproduces GPUMD's `c=1` q_scaler; `gpumd_init_parameters` re-inits coeffs **and** NN weights uniform(−1,1); L2 (`lambda_2`) shrinks the weights. |
-| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch) and is off by default. |
+| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch), is off by default, is per-stage (a stage-1 plateau jumps into Stage 2, surviving resume). |
+| `test_stream_mode.py` | `StreamDataStore` (the training data store): `collate` is bit-exact vs an independently assembled reference (concatenation + offsets + basis straight from the ops functions); metadata/mask consistency incl. missing channels; prefetched `iter_collated` matches direct collate; 2-rank DDP same-seed reproducibility. The DDP case is local-only: set `TORCHNEP_TEST_DDP=1` (skipped in CI — multi-process rendezvous is unreliable on shared runners). |
+| `test_compiled_autograd.py` | `CompiledAutogradForce` (make_fx-materialized autograd forces): outputs and parameter gradients (second-order path through the force loss) match eager autograd across batch shapes on one dynamic graph; energy-only calls fall back to eager. CUDA-only — auto-skipped on CPU hosts/CI. |
**Tolerance vs GPUMD:** `rtol=1e-5, atol=2e-4`.
diff --git a/tests/test_b1_and_gpumd_qscaler.py b/tests/test_b1_and_gpumd_qscaler.py
index 05245d7..661ba72 100644
--- a/tests/test_b1_and_gpumd_qscaler.py
+++ b/tests/test_b1_and_gpumd_qscaler.py
@@ -25,7 +25,7 @@
import torch
from torchnep.data import read_xyz, parse_nep_in
-from torchnep.train import (preprocess_structures, GPUDataStore,
+from torchnep.train import (preprocess_structures, StreamDataStore,
compute_q_scaler, recompute_b1_shift, train_nep)
from torchnep.model import NEPModel, gpumd_init_parameters
from _common import DATA_DIR
@@ -40,7 +40,7 @@
def _store(cfg, n=20, dtype=torch.float64):
frames = read_xyz(str(PBTE))[:n]
structs = preprocess_structures(frames, cfg, np.float64)
- return GPUDataStore(structs, torch.device("cpu"), dtype, config=cfg)
+ return StreamDataStore(structs, torch.device("cpu"), dtype, config=cfg)
def _cfg_from_nepin(tmp_path=None):
diff --git a/tests/test_compiled_autograd.py b/tests/test_compiled_autograd.py
new file mode 100644
index 0000000..65ff4a2
--- /dev/null
+++ b/tests/test_compiled_autograd.py
@@ -0,0 +1,129 @@
+# Copyright 2025 Yongchao Wu
+# This file is part of the TorchNEP project.
+# TorchNEP is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+# TorchNEP is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+# You should have received a copy of the GNU General Public License
+# along with TorchNEP. If not, see .
+
+"""Tests for CompiledAutogradForce (make_fx-materialized autograd forces).
+
+CUDA-only (the Inductor/Triton pipeline is the point) — auto-skipped on
+CPU-only hosts, so CI never runs the (slow) compile.
+"""
+import numpy as np
+import pytest
+import torch
+
+from torchnep.data import read_xyz, parse_nep_in
+from torchnep.train import StreamDataStore, preprocess_structures
+from torchnep.model import NEPModel, gpumd_init_parameters
+from _common import DATA_DIR
+
+PBTE = DATA_DIR.parent.parent / "example" / "PbTe" / "train.xyz"
+NEP_IN = ("type 2 Te Pb\ncutoff 6 4\nn_max 4 4\n"
+ "basis_size 6 6\nl_max 4 2 1\nneuron 30\n")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(),
+ reason="CompiledAutogradForce needs CUDA")
+
+
+def _setup(tmp_path, n_frames=24):
+ from torchnep.compiled_autograd import CompiledAutogradForce
+ dev = torch.device("cuda")
+ dtype = torch.float32
+ p = tmp_path / "nep.in"
+ p.write_text(NEP_IN)
+ cfg = parse_nep_in(str(p))
+ frames = read_xyz(str(PBTE))[:n_frames]
+ structs = preprocess_structures(frames, cfg, np.float32)
+ store = StreamDataStore(structs, dev, dtype, config=cfg)
+ torch.manual_seed(11)
+ model = NEPModel(cfg).to(dtype).to(dev)
+ gpumd_init_parameters(model)
+ model.set_q_scaler(torch.zeros(model.dim, device=dev),
+ torch.ones(model.dim, device=dev))
+ model.train()
+ return model, store, CompiledAutogradForce(model)
+
+
+def _run(fn, batch):
+ return fn(batch["rij_rad"], batch["rij_ang"],
+ batch["pair_i_rad"], batch["pair_j_rad"],
+ batch["pair_i_ang"], batch["pair_j_ang"],
+ batch["atom_types"], batch["N"],
+ batch["struct_idx"], batch["num_structures"],
+ need_forces=True, need_virial=True, backend="bmm")
+
+
+def _loss(r, batch):
+ e = ((r["Etot"] / batch["natoms"]
+ - batch["energy"] / batch["natoms"]) ** 2).mean()
+ f = ((r["forces"] - batch["forces"]) ** 2).mean()
+ v = (r["virial"] ** 2).mean()
+ return 0.01 * e + f + 0.01 * v
+
+
+def _rel(a, b):
+ return (a - b).abs().max().item() / max(a.abs().max().item(), 1e-30)
+
+
+def test_compiled_matches_eager_outputs_and_grads(tmp_path):
+ """One dynamic graph, multiple batch shapes: outputs AND parameter
+ gradients (second-order path through the force loss) must match the
+ eager autograd reference to float32 kernel-fusion noise."""
+ model, store, caf = _setup(tmp_path)
+ rng = np.random.default_rng(1)
+
+ for idx in (list(range(8)), rng.permutation(24)[:8].tolist(),
+ rng.permutation(24)[:20].tolist()):
+ batch = store.collate(idx)
+
+ model.zero_grad(set_to_none=True)
+ r_e = _run(model.compute_properties, batch)
+ _loss(r_e, batch).backward()
+ g_e = {n: p.grad.detach().clone()
+ for n, p in model.named_parameters() if p.grad is not None}
+
+ model.zero_grad(set_to_none=True)
+ r_c = _run(caf.compute_properties, batch)
+ _loss(r_c, batch).backward()
+ g_c = {n: p.grad.detach().clone()
+ for n, p in model.named_parameters() if p.grad is not None}
+
+ for key in ("Ei", "Etot", "forces", "virial"):
+ assert _rel(r_e[key].detach(), r_c[key].detach()) < 1e-4, key
+ assert set(g_e) == set(g_c)
+ for n in g_e:
+ denom = max(g_e[n].abs().max().item(), 1e-12)
+ assert (g_e[n] - g_c[n]).abs().max().item() / denom < 1e-3, n
+
+
+def test_compiled_energy_only_falls_back(tmp_path):
+ """need_forces=False routes to the eager model (graph is force-shaped)."""
+ model, store, caf = _setup(tmp_path, n_frames=8)
+ batch = store.collate(list(range(8)))
+ r = caf.compute_properties(
+ batch["rij_rad"], batch["rij_ang"],
+ batch["pair_i_rad"], batch["pair_j_rad"],
+ batch["pair_i_ang"], batch["pair_j_ang"],
+ batch["atom_types"], batch["N"],
+ batch["struct_idx"], batch["num_structures"],
+ need_forces=False, backend="bmm")
+ r_ref = model.compute_properties(
+ batch["rij_rad"], batch["rij_ang"],
+ batch["pair_i_rad"], batch["pair_j_rad"],
+ batch["pair_i_ang"], batch["pair_j_ang"],
+ batch["atom_types"], batch["N"],
+ batch["struct_idx"], batch["num_structures"],
+ need_forces=False, backend="bmm")
+ # Not torch.equal: CUDA scatter_add atomics make even two identical
+ # eager calls differ in the last ULP.
+ torch.testing.assert_close(r["Etot"], r_ref["Etot"],
+ rtol=1e-6, atol=1e-4)
+ assert "forces" not in r
diff --git a/tests/test_run_seed_and_valid.py b/tests/test_run_seed_and_valid.py
index 0c269f5..06f53e8 100644
--- a/tests/test_run_seed_and_valid.py
+++ b/tests/test_run_seed_and_valid.py
@@ -23,7 +23,7 @@
import torch
from torchnep.data import read_xyz, parse_nep_in
-from torchnep.train import (preprocess_structures, GPUDataStore, train_nep)
+from torchnep.train import (preprocess_structures, StreamDataStore, train_nep)
from torchnep.model import NEPModel
from _common import DATA_DIR
@@ -70,7 +70,7 @@ def _expected_split(n_frames, seed, ratio):
def _valid_energy_mse(cfg, frames, nep_txt):
"""Energy MSE of a saved model on the given frames (per-atom, eV/atom)."""
structs = preprocess_structures(frames, cfg, np.float64)
- ds = GPUDataStore(structs, torch.device("cpu"), torch.float64, config=cfg)
+ ds = StreamDataStore(structs, torch.device("cpu"), torch.float64, config=cfg)
m = NEPModel(cfg).to(torch.float64)
m.load_weights_from_nep_txt(nep_txt)
sq, n = 0.0, 0
diff --git a/tests/test_stream_mode.py b/tests/test_stream_mode.py
index a474fa8..1cb216a 100644
--- a/tests/test_stream_mode.py
+++ b/tests/test_stream_mode.py
@@ -11,190 +11,148 @@
# You should have received a copy of the GNU General Public License
# along with TorchNEP. If not, see .
-"""Tests for ``stream_mode`` (host-resident dataset, batch streaming).
+"""Tests for ``StreamDataStore`` — the (only) training data store.
-``StreamDataStore`` must be a drop-in replacement for ``GPUDataStore``:
-same collate values (bit-exact, including the on-the-fly basis), same
-metadata interface, and a training run in stream mode must reproduce the
-default-mode run exactly (CPU float64, where the backward is deterministic).
+The store keeps everything in host memory and assembles device batches on
+demand. ``collate`` must reproduce, bit-for-bit, a reference batch built
+independently from the raw per-frame structures (concatenation + index
+offsets + the Chebyshev/angular basis evaluated directly with the ops
+functions on the batch's rij).
"""
import numpy as np
import pytest
import torch
+from torchnep import ops
from torchnep.data import read_xyz, parse_nep_in
-from torchnep.train import (preprocess_structures, GPUDataStore,
- StreamDataStore, train_nep)
+from torchnep.train import (preprocess_structures, StreamDataStore,
+ iter_collated)
from _common import DATA_DIR, devices
PBTE = DATA_DIR.parent.parent / "example" / "PbTe" / "train.xyz"
NEP_IN = ("type 2 Te Pb\ncutoff 6 4\nn_max 4 4\n"
"basis_size 6 6\nl_max 4 2 1\nneuron 30\n")
-BATCH_KEYS = [
- "atom_types", "struct_idx", "pair_i_rad", "pair_j_rad", "rij_rad",
- "pair_i_ang", "pair_j_ang", "rij_ang", "energy", "natoms", "volumes",
- "energy_mask", "forces", "force_mask", "virial", "virial_mask",
- "fk_rad", "fkp_rad", "d12inv_rad", "fk_ang", "fkp_ang", "d12inv_ang",
- "blm",
-]
-
-def _config(tmp_path, extra=""):
+def _config(tmp_path):
p = tmp_path / "nep.in"
- p.write_text(NEP_IN + extra)
+ p.write_text(NEP_IN)
return parse_nep_in(str(p))
+def _reference_batch(structs, indices, cfg, dev, dtype):
+ """Independently assemble the batch straight from the structures."""
+ sel = [structs[i] for i in indices]
+ offsets = np.concatenate([[0], np.cumsum([s["natoms"] for s in sel])])
+
+ def cat(key):
+ return torch.from_numpy(np.concatenate([s[key] for s in sel])).to(dev)
+
+ ref = {
+ "N": int(offsets[-1]), "num_structures": len(sel),
+ "atom_types": cat("atom_types"),
+ "rij_rad": cat("rij_rad").to(dtype),
+ "rij_ang": cat("rij_ang").to(dtype),
+ }
+ for key, cnt in (("pair_i_rad", "rij_rad"), ("pair_j_rad", "rij_rad"),
+ ("pair_i_ang", "rij_ang"), ("pair_j_ang", "rij_ang")):
+ parts = [torch.from_numpy(s[key]) + int(offsets[k])
+ for k, s in enumerate(sel)]
+ ref[key] = torch.cat(parts).to(dev)
+ ref["struct_idx"] = torch.cat([
+ torch.full((s["natoms"],), k, dtype=torch.long)
+ for k, s in enumerate(sel)]).to(dev)
+
+ # Basis straight from the ops functions on the batch rij.
+ l3 = cfg["l_max"][0]
+ dr = torch.norm(ref["rij_rad"], dim=-1)
+ fk_r, fkp_r = ops.chebyshev_basis_and_deriv(
+ dr, cfg["cutoff_radial"], cfg["basis_size_radial"])
+ da = torch.norm(ref["rij_ang"], dim=-1)
+ fk_a, fkp_a = ops.chebyshev_basis_and_deriv(
+ da, cfg["cutoff_angular"], cfg["basis_size_angular"])
+ dinv_a = 1.0 / da
+ ref.update(fk_rad=fk_r, fkp_rad=fkp_r, d12inv_rad=1.0 / dr,
+ fk_ang=fk_a, fkp_ang=fkp_a, d12inv_ang=dinv_a,
+ blm=ops.angular_basis(
+ ref["rij_ang"][:, 0] * dinv_a,
+ ref["rij_ang"][:, 1] * dinv_a,
+ ref["rij_ang"][:, 2] * dinv_a, l3))
+ return ref
+
+
@pytest.mark.parametrize("device", devices())
-def test_collate_bit_exact(tmp_path, device):
- """Stream collate must reproduce GPUDataStore's batches bit-for-bit,
- including the recomputed Chebyshev/angular basis."""
+def test_collate_matches_reference(tmp_path, device):
+ """collate output must equal the independently assembled reference
+ bit-for-bit (indices, rij, and the on-the-fly basis)."""
dev = torch.device(device)
dtype = torch.float32
cfg = _config(tmp_path)
frames = read_xyz(str(PBTE))[:20]
structs = preprocess_structures(frames, cfg, np.float32)
-
- gpu_store = GPUDataStore(structs, dev, dtype, config=cfg)
- str_store = StreamDataStore(structs, dev, dtype, config=cfg)
+ store = StreamDataStore(structs, dev, dtype, config=cfg)
rng = np.random.default_rng(3)
for idx in ([0], list(range(8)), rng.permutation(20)[:8].tolist(),
list(range(20))):
- bg = gpu_store.collate(idx)
- bs = str_store.collate(idx)
- assert bg["N"] == bs["N"]
- assert bg["num_structures"] == bs["num_structures"]
- for key in BATCH_KEYS:
- tg, ts = bg[key], bs[key]
- assert tg.shape == ts.shape, f"{key}: shape mismatch"
- assert tg.dtype == ts.dtype, f"{key}: dtype mismatch"
- assert tg.device.type == ts.device.type, f"{key}: device mismatch"
- assert torch.equal(tg, ts), f"{key}: values differ"
+ batch = store.collate(idx)
+ ref = _reference_batch(structs, idx, cfg, dev, dtype)
+ assert batch["N"] == ref["N"]
+ assert batch["num_structures"] == ref["num_structures"]
+ for key in ("atom_types", "struct_idx", "pair_i_rad", "pair_j_rad",
+ "rij_rad", "pair_i_ang", "pair_j_ang", "rij_ang",
+ "fk_rad", "fkp_rad", "d12inv_rad", "fk_ang", "fkp_ang",
+ "d12inv_ang", "blm"):
+ assert torch.equal(batch[key], ref[key]), key
@pytest.mark.parametrize("device", devices())
-def test_store_metadata_parity(tmp_path, device):
- """Every metadata attribute consumers read must match GPUDataStore."""
+def test_store_metadata_and_masks(tmp_path, device):
+ """Metadata (counts / flags / per-frame views) and batch masks must
+ reflect the structures, including missing energy/forces channels."""
dev = torch.device(device)
cfg = _config(tmp_path)
frames = read_xyz(str(PBTE))[:12]
- # Exercise the missing-channel paths too (PbTe frames carry no virial,
- # so the virial-missing path is exercised by every frame).
frames[3].pop("energy", None)
frames[5].pop("forces", None)
- frames[7].pop("virial", None)
structs = preprocess_structures(frames, cfg, np.float64)
-
- a = GPUDataStore(structs, dev, torch.float64, config=cfg)
- b = StreamDataStore(structs, dev, torch.float64, config=cfg)
-
- assert a.n == b.n
- assert a.natoms == b.natoms
- assert a.energy == b.energy
- assert a.has_energy_flag == b.has_energy_flag
- assert a.has_forces_flag == b.has_forces_flag
- assert a.has_virial_flag == b.has_virial_flag
- assert (a.n_energy, a.n_forces, a.n_virial) == \
- (b.n_energy, b.n_forces, b.n_virial)
- assert (a.has_forces, a.has_virial) == (b.has_forces, b.has_virial)
- assert a.has_cached_basis and b.has_cached_basis
- assert torch.equal(a.volumes, b.volumes)
- for i in range(a.n):
- assert torch.equal(a.forces[i].cpu().double(),
- b.forces[i].cpu().double())
- assert torch.equal(a.virial[i].cpu().double(),
- b.virial[i].cpu().double())
- # Masks in a mixed-coverage batch
- ba = a.collate(list(range(12)))
- bb = b.collate(list(range(12)))
- for key in ("energy_mask", "force_mask", "virial_mask"):
- assert torch.equal(ba[key], bb[key])
-
-
-def _write_run_files(tmp_path, n_frames=20, epochs=3):
- nepin = tmp_path / "nep.in"
- nepin.write_text(NEP_IN + f"epoch {epochs}\nbatch 8\n")
- xyz = tmp_path / "train.xyz"
- raw = PBTE.read_text().splitlines()
- out, i, k = [], 0, 0
- while i < len(raw) and k < n_frames:
- na = int(raw[i].strip())
- out += raw[i:i + na + 2]
- i += na + 2; k += 1
- xyz.write_text("\n".join(out) + "\n")
- return str(nepin), str(xyz)
-
-
-def _train(nepin, xyz, out, **kw):
- kw.setdefault("device", "cpu")
- kw.setdefault("precision", "float64")
- kw.setdefault("print_interval", 100)
- kw.setdefault("restart", False)
- kw.setdefault("checkpoint_interval", 10000)
- kw.setdefault("prediction_interval", 10000)
- train_nep(config_file=nepin, data_file=xyz, output_dir=str(out), **kw)
-
-
-def _assert_files_numerically_equal(pa, pb, rtol=1e-8, atol=1e-12):
- """Token-wise comparison: text tokens must match exactly, numeric tokens
- within (rtol, atol).
-
- Byte-for-byte equality would be the ideal assertion (and holds on most
- machines), but the training forward passes go through BLAS matmuls whose
- reduction order can depend on heap alignment (MKL/OpenBLAS "conditional
- reproducibility") — on some CI hosts that costs ~1 ULP between two runs
- in the same process. The streamed inputs themselves ARE bit-exact (see
- test_collate_bit_exact, which is pure elementwise math and stays a strict
- torch.equal); any real divergence would exceed these tolerances by orders
- of magnitude.
- """
- import math
- ta, tb = pa.read_text().split(), pb.read_text().split()
- assert len(ta) == len(tb), f"{pa.name}: token count {len(ta)} != {len(tb)}"
- for k, (x, y) in enumerate(zip(ta, tb)):
- try:
- fx, fy = float(x), float(y)
- except ValueError:
- assert x == y, f"{pa.name} token {k}: {x!r} != {y!r}"
- continue
- assert math.isclose(fx, fy, rel_tol=rtol, abs_tol=atol), \
- f"{pa.name} token {k}: {x} vs {y}"
-
-
-def test_train_stream_reproduces_default(tmp_path):
- """Same seed, stream_mode on vs off: matching loss.out, nep_final.txt
- and end-of-training predictions (the streamed batches are bit-identical;
- outputs compared numerically to tolerate BLAS-alignment ULP noise on
- some CI hosts — see _assert_files_numerically_equal)."""
- nepin, xyz = _write_run_files(tmp_path, n_frames=16, epochs=3)
-
- out_a = tmp_path / "out_default"
- out_b = tmp_path / "out_stream"
- _train(nepin, xyz, out_a, run_seed=77, prediction_interval=2)
- _train(nepin, xyz, out_b, run_seed=77, prediction_interval=2,
- stream_mode=True)
-
- for f in ("loss.out", "nep_final.txt", "nep_best.txt",
- "energy_train.out", "force_train.out", "virial_train.out",
- "stress_train.out"):
- _assert_files_numerically_equal(out_a / f, out_b / f)
-
-
-def test_train_stream_with_validation(tmp_path):
- """stream_mode combined with valid_ratio: the validation store streams
- too, and the run matches the default-mode run exactly."""
- nepin, xyz = _write_run_files(tmp_path, n_frames=20, epochs=3)
-
- out_a = tmp_path / "out_default"
- out_b = tmp_path / "out_stream"
- _train(nepin, xyz, out_a, run_seed=5, valid_ratio=0.25)
- _train(nepin, xyz, out_b, run_seed=5, valid_ratio=0.25, stream_mode=True)
-
- for f in ("loss.out", "nep_best.txt",
- "energy_test.out", "force_test.out", "virial_test.out"):
- _assert_files_numerically_equal(out_a / f, out_b / f)
+ store = StreamDataStore(structs, dev, torch.float64, config=cfg)
+
+ assert store.n == 12
+ assert store.natoms == [s["natoms"] for s in structs]
+ assert store.has_energy_flag == ["energy" in s for s in structs]
+ assert store.has_forces_flag == ["forces" in s for s in structs]
+ assert store.n_energy == 11 and store.n_forces == 11
+ for i in (0, 3, 5, 11):
+ exp = structs[i].get("forces")
+ if exp is not None:
+ assert torch.equal(store.forces[i].cpu().double(),
+ torch.from_numpy(np.asarray(exp)).double())
+
+ batch = store.collate(list(range(12)))
+ e_mask = batch["energy_mask"].cpu().tolist()
+ assert e_mask == store.has_energy_flag
+ f_mask = batch["force_mask"].cpu()
+ off = np.concatenate([[0], np.cumsum(store.natoms)])
+ for i in range(12):
+ seg = f_mask[int(off[i]):int(off[i + 1])]
+ assert bool(seg.all()) == store.has_forces_flag[i]
+ assert bool(seg.any()) == store.has_forces_flag[i]
+
+
+def test_iter_collated_prefetch_matches_direct(tmp_path):
+ """Prefetched iteration yields exactly the same batches as direct
+ collate calls (same order, same tensors)."""
+ cfg = _config(tmp_path)
+ frames = read_xyz(str(PBTE))[:16]
+ structs = preprocess_structures(frames, cfg, np.float64)
+ store = StreamDataStore(structs, torch.device("cpu"), torch.float64,
+ config=cfg)
+ idx_lists = [[0, 3, 5], [1, 2], list(range(16)), [15]]
+ direct = [store.collate(i) for i in idx_lists]
+ for got, want in zip(iter_collated(store, idx_lists), direct):
+ for key in ("atom_types", "rij_rad", "fk_ang", "forces", "energy"):
+ assert torch.equal(got[key], want[key]), key
_SHARDED_RUNNER = """
@@ -203,15 +161,13 @@ def test_train_stream_with_validation(tmp_path):
train_nep_sharded(sys.argv[1], sys.argv[2], output_dir=sys.argv[3],
precision="float64", print_interval=100,
checkpoint_interval=10000, prediction_interval=10000,
- restart=False, run_seed=99,
- stream_mode=(sys.argv[4] == "stream"))
+ restart=False, run_seed=99)
"""
-def test_sharded_stream_matches_default(tmp_path):
- """2-rank DDP (CPU/gloo): stream_mode reproduces the default sharded
- run (numeric comparison, same tolerance rationale as the single-GPU
- reproduction test).
+def test_sharded_run_reproducible(tmp_path):
+ """2-rank DDP (CPU/gloo): two identical runs with the same seed are
+ byte-identical — DDP smoke coverage for the streamed store.
Opt-in (local only): multi-process rendezvous can hang on constrained
CI runners, so this is skipped unless TORCHNEP_TEST_DDP=1 is set.
@@ -226,22 +182,29 @@ def test_sharded_stream_matches_default(tmp_path):
if torchrun is None:
pytest.skip("torchrun not on PATH")
- nepin, xyz = _write_run_files(tmp_path, n_frames=16, epochs=3)
+ nepin = tmp_path / "nep.in"
+ nepin.write_text(NEP_IN + "epoch 3\nbatch 4\n")
+ raw = PBTE.read_text().splitlines()
+ out, i, k = [], 0, 0
+ while i < len(raw) and k < 16:
+ na = int(raw[i].strip())
+ out += raw[i:i + na + 2]
+ i += na + 2; k += 1
+ xyz = tmp_path / "train.xyz"
+ xyz.write_text("\n".join(out) + "\n")
runner = tmp_path / "runner.py"
runner.write_text(_SHARDED_RUNNER)
- import os
root = str(DATA_DIR.parent.parent)
env = dict(os.environ, CUDA_VISIBLE_DEVICES="",
PYTHONPATH=root + os.pathsep + os.environ.get("PYTHONPATH", ""))
- for out, mode in (("out_default", "default"), ("out_stream", "stream")):
+ for out_dir in ("out_a", "out_b"):
r = subprocess.run(
[torchrun, "--standalone", "--nproc_per_node=2", str(runner),
- nepin, xyz, str(tmp_path / out), mode],
+ str(nepin), str(xyz), str(tmp_path / out_dir)],
capture_output=True, text=True, env=env, timeout=600)
assert r.returncode == 0, r.stderr[-2000:]
- a, b = tmp_path / "out_default", tmp_path / "out_stream"
- for f in ("loss.out", "nep_best.txt",
- "energy_train.out", "force_train.out", "virial_train.out"):
- _assert_files_numerically_equal(a / f, b / f)
+ a, b = tmp_path / "out_a", tmp_path / "out_b"
+ assert (a / "loss.out").read_text() == (b / "loss.out").read_text()
+ assert (a / "nep_final.txt").read_text() == (b / "nep_final.txt").read_text()
diff --git a/torchnep/compiled_autograd.py b/torchnep/compiled_autograd.py
new file mode 100644
index 0000000..af9d3b2
--- /dev/null
+++ b/torchnep/compiled_autograd.py
@@ -0,0 +1,315 @@
+# Copyright 2025 Yongchao Wu
+# This file is part of the TorchNEP project.
+# TorchNEP is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+# TorchNEP is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+# You should have received a copy of the GNU General Public License
+# along with TorchNEP. If not, see .
+
+"""torch.compile for the AUTOGRAD force path via make_fx.
+
+The autograd path computes forces as F = -dE/drij by a nested
+``autograd.grad(create_graph=True)``, which torch.compile cannot lower
+directly — that is why ``use_compile`` used to be ignored for
+``use_autograd_forces=True``. The workaround (the approach follows the idea
+used by DeepMD-kit's compiled models):
+
+ 1. write (energy, forces) as a PURE function of (params+buffers, batch
+ geometry) — parameters enter as function inputs, so a normal
+ ``loss.backward()`` still fills ``model.parameters().grad``;
+ 2. ``make_fx`` trace it symbolically: the inner first-order gradient is
+ materialized as ordinary FX ops (no runtime double-backward left);
+ 3. ``strip_detach``: make_fx wraps saved activations in aten.detach chains
+ that would silently CUT the second-order path from the force loss back
+ to the parameters (param grads come out wrong, rel err ~1) — the traced
+ leaves are already detached, so every detach node is safe to drop;
+ 4. ``torch.compile`` the resulting graph with dynamic shapes, so ONE graph
+ serves every batch size.
+
+The traced energy re-implements the per-type NN dispatch with
+``torch.where`` over all types (NEPModel.forward's ``mask.any()`` branches
+are data-dependent and cannot be traced symbolically), and uses the "bmm"
+contraction backend (the "loop" backend's per-type-pair ``.any()`` masks are
+data-dependent too). ZBL stays OUTSIDE the compiled graph (its typewise
+cutoffs use ``.item()``) and is added eagerly, exactly like the cached
+analytical path does.
+"""
+
+import os
+import torch
+from torch.fx.experimental.proxy_tensor import make_fx
+
+from . import ops
+
+
+_PATCHED = False
+
+
+def apply_compile_patches():
+ """Process-global adjustments for a robust dynamic-shape compile.
+
+ Idempotent. The key one is disabling opt_einsum: it flattens contraction
+ operands and bakes the pair/atom count into a constant, forcing a
+ recompile for EVERY distinct batch shape. The einsums here contract tiny
+ angular dims, so the naive path costs nothing at runtime.
+ """
+ global _PATCHED
+ if _PATCHED:
+ return
+ os.environ.setdefault("TORCHINDUCTOR_FX_GRAPH_CACHE", "1")
+ torch.backends.opt_einsum.enabled = False
+ try:
+ torch._dynamo.config.recompile_limit = 64
+ torch._dynamo.config.accumulated_recompile_limit = 256
+ except Exception:
+ pass
+ _patch_cantsplit()
+ _PATCHED = True
+
+
+def _patch_cantsplit():
+ """torch 2.11: the Triton tiler does ``raise CantSplit`` with no args while
+ CantSplit.__init__ requires them — a TypeError escapes instead of the
+ catchable control-flow exception. Give it an arg-free constructor."""
+ try:
+ from torch._inductor.codegen.simd import CantSplit
+ except Exception:
+ try:
+ from torch._inductor.exc import CantSplit
+ except Exception:
+ return
+ if getattr(CantSplit, "_torchnep_patched", False):
+ return
+
+ def __init__(self, expr=None, remaining=None):
+ Exception.__init__(self)
+ self.expr, self.remaining = expr, remaining
+
+ CantSplit.__init__ = __init__
+ CantSplit._torchnep_patched = True
+
+
+def strip_detach(gm):
+ """Remove every aten.detach node from a make_fx graph (see module doc)."""
+ det = torch.ops.aten.detach.default
+ for n in list(gm.graph.nodes):
+ if n.op == "call_function" and n.target == det:
+ n.replace_all_uses_with(n.args[0])
+ gm.graph.erase_node(n)
+ gm.graph.lint()
+ gm.recompile()
+ return gm
+
+
+def inductor_options():
+ """Conservative Inductor options: the default pipeline runs heavy
+ fusion/autotune on the large
+ materialized second-order graph — minutes of codegen, and small GPUs
+ can't even run max_autotune_gemm. This cuts compile to seconds."""
+ opts = {
+ "max_autotune": False,
+ "shape_padding": True,
+ "epilogue_fusion": False,
+ "triton.cudagraphs": False,
+ "max_fusion_size": 8,
+ "triton.persistent_reductions": False,
+ "triton.mix_order_reduction": False,
+ "triton.max_tiles": 1,
+ }
+ try:
+ from torch._inductor import config as ic
+ valid = ic.get_config_copy()
+ opts = {k: v for k, v in opts.items() if k.replace("-", "_") in valid}
+ except Exception:
+ pass
+ return opts
+
+
+class CompiledAutogradForce:
+ """Lazily make_fx-traces + compiles the autograd force on first use, then
+ one dynamic-shape graph serves all batch sizes.
+
+ ``compute_properties`` mirrors ``NEPModel.compute_properties`` (same
+ signature, same result dict), so the training/eval loops can use it as a
+ drop-in replacement for the eager autograd path. Trace lazily AFTER
+ q_scaler is set — buffers are graph inputs, but tracing needs the model
+ fully initialized.
+ """
+
+ def __init__(self, model):
+ apply_compile_patches()
+ self.model = model
+ pd = {**dict(model.named_parameters()), **dict(model.named_buffers())}
+ self.pkeys = list(pd)
+ self._compiled = None
+
+ def _pvals(self):
+ d = {**dict(self.model.named_parameters()),
+ **dict(self.model.named_buffers())}
+ return tuple(d[k] for k in self.pkeys)
+
+ # -- pure traced function -------------------------------------------------
+
+ def _energy(self, pd, rij_rad, rij_ang, pi_rad, pj_rad, pi_ang, pj_ang,
+ atom_types):
+ """Per-atom NN energy as a pure function of (params, geometry).
+
+ Re-implements NEPModel.forward without data-dependent branches:
+ every type's net runs on all atoms, torch.where selects. All
+ parameters appear in the graph, so DDP-style full-graph gradient
+ coverage is preserved too.
+ """
+ m = self.model
+ N = atom_types.shape[0]
+ q = ops.compute_descriptors(
+ rij_rad, rij_ang, pi_rad, pj_rad, pi_ang, pj_ang,
+ atom_types, N, pd["c_param_2"], pd.get("c_param_3"),
+ m.rc_radial, m.rc_angular,
+ m.basis_size_radial, m.basis_size_angular,
+ m.n_max_radial, m.n_max_angular,
+ m.l_max_3b,
+ m.has_q_222, m.has_q_1111, m.has_q_112,
+ m.num_lm, pd["_c3b"], pd["_c4b"], pd["_c5b"],
+ pd["_c4b2"],
+ rij_rad.dtype, rij_rad.device,
+ backend="bmm",
+ has_q_123=m.has_q_123, has_q_233=m.has_q_233,
+ has_q_134=m.has_q_134,
+ )
+ q_scaled = q * pd["q_scaler"]
+ Ei = torch.zeros(N, dtype=q.dtype, device=q.device)
+ for t in range(m.num_types):
+ w0 = pd[f"fitting_nets.{t}.w0"]
+ b0 = pd[f"fitting_nets.{t}.b0"]
+ w1 = pd[f"fitting_nets.{t}.w1"]
+ e_t = torch.tanh(q_scaled @ w0 - b0) @ w1
+ Ei = Ei + torch.where(atom_types == t, e_t,
+ torch.zeros((), dtype=q.dtype,
+ device=q.device))
+ return Ei - pd["b1"]
+
+ def _raw(self, pvals, rij_rad, rij_ang, pi_rad, pj_rad, pi_ang, pj_ang,
+ atom_types):
+ pd = dict(zip(self.pkeys, pvals))
+ N = atom_types.shape[0]
+ rr = rij_rad.detach().requires_grad_(True)
+ ra = rij_ang.detach().requires_grad_(True)
+ Ei = self._energy(pd, rr, ra, pi_rad, pj_rad, pi_ang, pj_ang,
+ atom_types)
+ gr, ga = torch.autograd.grad(Ei.sum(), [rr, ra], create_graph=True,
+ allow_unused=True)
+ if gr is None:
+ gr = torch.zeros_like(rr)
+ if ga is None:
+ ga = torch.zeros_like(ra)
+ forces, virial = ops.accumulate_forces_virial(
+ N, pi_rad, pj_rad, rr, gr, pi_ang, pj_ang, ra, ga,
+ rr.dtype, rr.device)
+ return Ei, forces, virial
+
+ # -- tracing --------------------------------------------------------------
+
+ def _prime_args(self, args):
+ """Synthetic tracing inputs with PRIME atom/pair counts: make_fx
+ then derives clean symbolic dims, so no batch dim can
+ accidentally equal a channel/basis count and get mis-specialized as
+ a constant."""
+ pvals = args[0]
+ dev = next(iter(pvals)).device if pvals else args[1].device
+ dtype = args[1].dtype
+ m = self.model
+ na, npr, npa = 53, 101, 89 # distinct primes >> any model dim
+ g = torch.Generator(device="cpu").manual_seed(0)
+
+ def _rij(n, rc):
+ v = torch.randn(n, 3, generator=g, dtype=torch.float64)
+ v = v / v.norm(dim=1, keepdim=True)
+ r = 0.5 + torch.rand(n, 1, generator=g, dtype=torch.float64) \
+ * (0.8 * rc - 0.5)
+ return (v * r).to(device=dev, dtype=dtype)
+
+ s_types = torch.arange(na, device=dev) % m.num_types
+ s_pir = torch.randint(0, na, (npr,), generator=g).to(dev)
+ s_pjr = torch.randint(0, na, (npr,), generator=g).to(dev)
+ s_pia = torch.randint(0, na, (npa,), generator=g).to(dev)
+ s_pja = torch.randint(0, na, (npa,), generator=g).to(dev)
+ return (pvals, _rij(npr, m.rc_radial), _rij(npa, m.rc_angular),
+ s_pir, s_pjr, s_pia, s_pja, s_types)
+
+ def _ensure_compiled(self, args):
+ if self._compiled is not None:
+ return
+ # Trace a plain closure — make_fx on the bound method would count
+ # ``self`` as an input.
+ fn = lambda *a: self._raw(*a)
+ gm = make_fx(fn, tracing_mode="symbolic",
+ _allow_non_fake_inputs=True)(*self._prime_args(args))
+ strip_detach(gm)
+ self._compiled = torch.compile(gm, dynamic=True,
+ options=inductor_options())
+
+ # -- public: drop-in for NEPModel.compute_properties ----------------------
+
+ def compute_properties(self, rij_rad, rij_ang, pi_rad, pj_rad,
+ pi_ang, pj_ang, atom_types, N,
+ struct_idx, num_structures,
+ need_forces=True, need_virial=False,
+ backend: str = "bmm"):
+ """Same signature/result dict as ``NEPModel.compute_properties``.
+
+ The compiled graph always evaluates forces+virial (the graph is
+ traced once); energy-only calls fall back to the eager model —
+ they are rare (b1/eval passes) and cheap.
+ """
+ if not need_forces:
+ return self.model.compute_properties(
+ rij_rad, rij_ang, pi_rad, pj_rad, pi_ang, pj_ang,
+ atom_types, N, struct_idx, num_structures,
+ need_forces=False, need_virial=need_virial, backend=backend)
+
+ args = (self._pvals(), rij_rad, rij_ang, pi_rad, pj_rad,
+ pi_ang, pj_ang, atom_types)
+ self._ensure_compiled(args)
+ Ei, forces, virial = self._compiled(*args)
+
+ m = self.model
+ dtype, device = rij_rad.dtype, rij_rad.device
+ if m.zbl is not None:
+ # ZBL outside the compiled graph (typewise cutoffs use .item()).
+ # No trainable parameters — energies/forces detach, same as the
+ # cached analytical path.
+ with torch.enable_grad():
+ rz = rij_ang.detach().requires_grad_(True)
+ Ei_zbl = ops.compute_zbl(
+ atom_types, pi_ang, pj_ang, rz, N,
+ m.atomic_numbers.tolist(),
+ m.zbl_rc_inner, m.zbl_rc_outer, m.zbl_typewise_factor,
+ getattr(m, "zbl_rc_inner_per_type", None),
+ getattr(m, "zbl_rc_outer_per_type", None), dtype, device)
+ if Ei_zbl.requires_grad:
+ g_zbl = torch.autograd.grad(Ei_zbl.sum(), rz,
+ allow_unused=True)[0]
+ else:
+ g_zbl = None
+ Ei = Ei + Ei_zbl.detach()
+ if g_zbl is not None:
+ empty_i = torch.zeros(0, dtype=torch.long, device=device)
+ empty_r = torch.zeros(0, 3, dtype=dtype, device=device)
+ zf, zv = ops.accumulate_forces_virial(
+ N, empty_i, empty_i, empty_r, empty_r,
+ pi_ang, pj_ang, rij_ang.detach(), g_zbl.detach(),
+ dtype, device)
+ forces = forces + zf
+ virial = virial + zv
+
+ Etot = torch.zeros(num_structures, dtype=dtype, device=device)
+ Etot.scatter_add_(0, struct_idx, Ei)
+ result = {"Ei": Ei, "Etot": Etot, "forces": forces}
+ if need_virial:
+ result["virial"] = virial
+ return result
diff --git a/torchnep/model.py b/torchnep/model.py
index f23bd0e..3829108 100644
--- a/torchnep/model.py
+++ b/torchnep/model.py
@@ -262,14 +262,21 @@ def compute_properties(self, rij_rad, rij_ang, pi_rad, pj_rad,
return result
- def compute_properties_cached(self, batch, need_forces=True, need_virial=False,
- backend: str = "loop"):
- """Compute energy, forces, virial using precomputed basis.
-
- Uses fully analytical force computation — no create_graph=True needed.
- Forces are differentiable through c2, c3 (via Fp->NN weights and via s->c3).
-
- ``backend`` in {"loop", "bmm"} — see torchnep.ops.resolve_backend.
+ def _cached_core(self, batch, need_forces=True, need_virial=False,
+ backend: str = "loop"):
+ """Descriptor + NN + analytical-force part of the cached compute.
+
+ Deliberately free of data-dependent Python control flow (the
+ per-type NN dispatch is branchless: every type's net runs on all
+ atoms and ``torch.where`` selects), and ZBL is NOT included — so
+ ``torch.compile`` captures this whole function as ONE graph with no
+ breaks (the eager path's ``mask.any()`` branches and the ZBL block's
+ inner ``autograd.grad`` each split the graph, leaving ~2.7x more
+ kernel launches). ZBL and result assembly live in the
+ ``compute_properties_cached`` wrapper.
+
+ Returns ``(Ei, forces, virial)`` — forces/virial are None when not
+ requested.
"""
dtype = self.q_scaler.dtype
device = self.q_scaler.device
@@ -316,46 +323,83 @@ def compute_properties_cached(self, batch, need_forces=True, need_virial=False,
# NN forward + Fp computation (differentiable through NN weights).
#
- # Every per-type fitting net is touched in the graph on every forward
- # — even types with no atoms in this batch get a zeroed-out dummy pass.
- # This keeps DDP gradient bookkeeping consistent (no need for
- # find_unused_parameters=True) and, critically, avoids the implicit
- # /world_size gradient dilution that DDP applies to unused parameters
- # (which was biasing rare-type NNs toward lower effective LR).
+ # Branchless per-type dispatch: EVERY type's net runs on all atoms
+ # and torch.where selects per atom. The nets are tiny, so the extra
+ # flops are negligible; in exchange there is no data-dependent
+ # Python branch (torch.compile keeps one graph) and every parameter
+ # is in the autograd graph on every forward — which also keeps DDP
+ # gradient bookkeeping consistent without find_unused_parameters
+ # and avoids the implicit /world_size gradient dilution for types
+ # absent from a batch.
Ei = torch.zeros(N, dtype=dtype, device=device)
Fp = torch.zeros(N, self.dim, dtype=dtype, device=device)
- dummy_accum = torch.zeros((), dtype=dtype, device=device)
- dummy_q = q_scaled[:1] if q_scaled.shape[0] > 0 else torch.zeros(
- 1, self.dim, dtype=dtype, device=device)
-
for t in range(self.num_types):
- mask = batch["atom_types"] == t
net = self.fitting_nets[t]
- if mask.any():
- qt = q_scaled[mask]
- z = qt @ net.w0 - net.b0
- h = torch.tanh(z)
- Ei[mask] = h @ net.w1
- tanh_der = 1.0 - h * h
- Fp[mask] = (net.w1 * tanh_der) @ net.w0.T
- else:
- # Dummy forward (the * 0 below nulls the contribution but
- # keeps the net's parameters in the autograd graph).
- z_d = dummy_q @ net.w0 - net.b0
- h_d = torch.tanh(z_d)
- dummy_accum = dummy_accum + (h_d @ net.w1).sum()
+ z = q_scaled @ net.w0 - net.b0
+ h = torch.tanh(z)
+ e_t = h @ net.w1
+ tanh_der = 1.0 - h * h
+ fp_t = (net.w1 * tanh_der) @ net.w0.T
+ sel = batch["atom_types"] == t
+ Ei = torch.where(sel, e_t, Ei)
+ Fp = torch.where(sel.unsqueeze(-1), fp_t, Fp)
Fp = Fp * self.q_scaler # absorb q_scaler into Fp
- # Nail the unused-type gradient path into Ei without changing its value.
- Ei = Ei + dummy_accum * 0.0
Ei = Ei - self.b1 # subtract shared output bias
+ forces = None
+ virial = None
+ if need_forces:
+ # Analytical forces: fully differentiable through c2/c3 and NN
+ # weights (Fp). No create_graph=True needed — chain rule is
+ # computed explicitly.
+ forces, virial = ops.compute_analytical_forces(
+ Fp, batch["atom_types"], N,
+ self.c_param_2, self.c_param_3,
+ batch["fkp_rad"], batch["fkp_ang"], batch["blm"],
+ batch["pair_i_rad"], batch["pair_j_rad"],
+ batch["rij_rad"], batch["d12inv_rad"],
+ batch["pair_i_ang"], batch["pair_j_ang"],
+ batch["rij_ang"], batch["d12inv_ang"],
+ s, gn_ang,
+ self.n_max_radial, self.n_max_angular,
+ self.l_max_3b,
+ self.has_q_222, self.has_q_1111, self.has_q_112,
+ self.num_lm, self._c3b, self._c4b, self._c5b,
+ self._c4b2,
+ dtype, device,
+ compute_virial=need_virial,
+ backend=backend,
+ has_q_123=self.has_q_123, has_q_233=self.has_q_233,
+ has_q_134=self.has_q_134,
+ )
+ return Ei, forces, virial
+
+ def compute_properties_cached(self, batch, need_forces=True, need_virial=False,
+ backend: str = "loop", core_fn=None):
+ """Compute energy, forces, virial using precomputed basis.
+
+ Uses fully analytical force computation — no create_graph=True needed.
+ Forces are differentiable through c2, c3 (via Fp->NN weights and via s->c3).
+
+ ``backend`` in {"loop", "bmm"} — see torchnep.ops.resolve_backend.
+ ``core_fn`` optionally substitutes a ``torch.compile``d version of
+ ``_cached_core`` (the trainer passes one); the ZBL add-on and result
+ assembly below stay eager either way (ZBL's typewise cutoffs and
+ inner autograd.grad cannot be captured in the compiled graph).
+ """
+ dtype = self.q_scaler.dtype
+ device = self.q_scaler.device
+ N = batch["N"]
+
+ core = core_fn if core_fn is not None else self._cached_core
+ Ei, forces, virial = core(batch, need_forces=need_forces,
+ need_virial=need_virial, backend=backend)
+
# ZBL energy + forces (no trainable params; local autograd on rij_ang).
# enable_grad: end-of-training predict_from_store wraps this call in
# torch.no_grad(), under which Ei_zbl.requires_grad would be False
# and the ZBL force contribution would be silently dropped.
- zbl_forces = None
- zbl_virial = None
if self.zbl is not None:
with torch.enable_grad():
rij_zbl = batch["rij_ang"].detach().requires_grad_(True)
@@ -380,43 +424,19 @@ def compute_properties_cached(self, batch, need_forces=True, need_virial=False,
batch["rij_ang"].detach(), g_zbl.detach(),
dtype, device,
)
+ if forces is not None:
+ forces = forces + zbl_forces
+ if need_virial and virial is not None:
+ virial = virial + zbl_virial
Etot = torch.zeros(batch["num_structures"], dtype=dtype, device=device)
Etot.scatter_add_(0, batch["struct_idx"], Ei)
result = {"Ei": Ei, "Etot": Etot}
-
if need_forces:
- # Analytical forces: fully differentiable through c2/c3 and NN weights (Fp).
- # No create_graph=True needed — chain rule is computed explicitly.
- forces, virial = ops.compute_analytical_forces(
- Fp, batch["atom_types"], N,
- self.c_param_2, self.c_param_3,
- batch["fkp_rad"], batch["fkp_ang"], batch["blm"],
- batch["pair_i_rad"], batch["pair_j_rad"],
- batch["rij_rad"], batch["d12inv_rad"],
- batch["pair_i_ang"], batch["pair_j_ang"],
- batch["rij_ang"], batch["d12inv_ang"],
- s, gn_ang,
- self.n_max_radial, self.n_max_angular,
- self.l_max_3b,
- self.has_q_222, self.has_q_1111, self.has_q_112,
- self.num_lm, self._c3b, self._c4b, self._c5b,
- self._c4b2,
- dtype, device,
- compute_virial=need_virial,
- backend=backend,
- has_q_123=self.has_q_123, has_q_233=self.has_q_233,
- has_q_134=self.has_q_134,
- )
- if zbl_forces is not None:
- forces = forces + zbl_forces
- if need_virial:
- virial = virial + zbl_virial
result["forces"] = forces
if need_virial and virial is not None:
result["virial"] = virial
-
return result
def load_weights_from_nep_txt(self, path: str):
diff --git a/torchnep/ops.py b/torchnep/ops.py
index 1486e5d..dfb088e 100644
--- a/torchnep/ops.py
+++ b/torchnep/ops.py
@@ -47,8 +47,14 @@ def resolve_backend(backend: str = "auto",
under torch.compile -> "bmm" (vectorised path fuses far better; the per-
type Python loop forces graph breaks, so bmm
is consistently fastest once compiled)
- ntypes >= 8 -> "bmm" (fancy-index + batched GEMM wins)
- otherwise -> "loop" (few-types eager; inline Python loop fastest)
+ ntypes >= 20 -> "bmm" (the O(ntypes^2) loop launches finally lose
+ to one batched GEMM)
+ otherwise -> "loop" (eager; inline Python loop fastest)
+
+ The eager threshold comes from a 6000-frame benchmark sweep (16-element
+ alloy set, A2000): loop beat bmm clearly up to 8 types (8.8 vs 14.1
+ s/epoch) and only reached parity at 16 (15.6 vs 15.1), so the crossover
+ sits near ~20 types.
Any non-"auto" string is returned unchanged (explicit override wins).
"""
@@ -56,7 +62,7 @@ def resolve_backend(backend: str = "auto",
return backend
if use_compile:
return "bmm"
- if num_types is not None and num_types >= 8:
+ if num_types is not None and num_types >= 20:
return "bmm"
return "loop"
diff --git a/torchnep/predict.py b/torchnep/predict.py
index 4f486cc..c15e2f2 100644
--- a/torchnep/predict.py
+++ b/torchnep/predict.py
@@ -393,7 +393,7 @@ def _log(msg):
# ---------------------------------------------------------------------------
-# End-of-training prediction that reuses the in-memory model + GPUDataStore
+# End-of-training prediction that reuses the in-memory model + data store
# (no xyz re-read, no neighbor-list rebuild, no second GPU upload).
# ---------------------------------------------------------------------------
@@ -402,7 +402,7 @@ def predict_from_store(model, data_store, output_dir: str,
backend: str = "auto",
verbose: bool = True,
suffix: str = "train"):
- """Run prediction using an already-loaded NEPModel + GPUDataStore.
+ """Run prediction using an already-loaded NEPModel + StreamDataStore.
Designed for the end of training: reuses the preprocessed data_store so
there is no xyz re-read / neighbor-list rebuild / GPU upload. The
@@ -638,7 +638,7 @@ def predict_from_store_sharded(model, data_store, local_global_idx,
Parameters
----------
model : NEPModel (DDP replica — parameters are in sync across ranks).
- data_store : GPUDataStore (this rank's local shard).
+ data_store : StreamDataStore (this rank's local shard).
local_global_idx : list[int] original xyz-frame index for each local
frame (length == ``data_store.n``). Supplied by the random shard
assignment in ``train_nep_sharded``.
diff --git a/torchnep/train.py b/torchnep/train.py
index 69a34c5..aeacd15 100644
--- a/torchnep/train.py
+++ b/torchnep/train.py
@@ -209,304 +209,25 @@ def tag(*keys):
# ---------------------------------------------------------------------------
-# GPU data store — all data pre-loaded to device
+# Data store — host-resident, batches streamed to the device
# ---------------------------------------------------------------------------
-def _basis_chunk_size(device, dtype, basis_size_angular, num_lm, l_max_3b,
- min_chunk=1 << 16, max_chunk=1 << 23):
- """Pairs-per-chunk for the cached-basis precompute in ``GPUDataStore``.
-
- The basis is built in chunks so the transient working set is one chunk
- instead of the whole shard — this lowers the construction-time GPU memory
- peak (which otherwise sits well above the steady training footprint and
- needlessly caps how big a shard each rank can hold). Results are unchanged:
- the Chebyshev/angular bases are per-pair elementwise, so chunking is
- bit-identical to the one-shot path.
-
- The chunk is sized so one chunk's transient stays within a small fixed
- budget (``TARGET`` below, ~512 MB) — the whole point is to keep the peak
- just above the steady footprint, not to go as fast as possible. On CUDA the
- budget is additionally clamped to a fraction of the memory still *free*
- after the persistent basis buffers are allocated (queried via
- ``mem_get_info``), so a memory-tight card shrinks the chunk further rather
- than OOM-ing. On CPU only the fixed budget applies. Bigger chunks make
- ``__init__`` faster; they never affect training speed (training reads the
- same split-view layout either way).
- """
- elem = 8 if dtype == torch.float64 else 4
- # Generous per-pair transient estimate: Chebyshev scratch (~14 vectors),
- # angular z/Re/Im powers + the blm list and its torch.stack copy
- # (~2*num_lm + 4*(l+1)), plus margin. Over-estimating only shrinks the
- # chunk, which is safe.
- per_pair = elem * (2 * (basis_size_angular + 1) + 2 * num_lm
- + 4 * (l_max_3b + 1) + 40)
- TARGET = 512 * 1024 * 1024 # ~512 MB transient budget per chunk
- budget = TARGET
- if device.type == "cuda":
- try:
- free, _ = torch.cuda.mem_get_info(device)
- # Never let the transient eat more than half of what's free, so a
- # tight card shrinks the chunk instead of OOM-ing at build time.
- budget = min(TARGET, int(free * 0.5))
- except Exception:
- budget = min(TARGET, 256 * 1024 * 1024)
- chunk = budget // max(per_pair, 1)
- return max(min_chunk, min(max_chunk, chunk))
-
-
-class GPUDataStore:
- """Pre-loads all structure data to GPU for zero-copy batch collation.
-
- When ``config`` is given, also caches Chebyshev basis functions and
- angular basis on GPU so training never recomputes them.
- """
-
- def __init__(self, structures: List[Dict], device: torch.device,
- dtype: torch.dtype, config: dict = None):
- self.device = device
- self.dtype = dtype
- self.n = len(structures)
- self.has_cached_basis = config is not None
-
- n_rad = np.array([len(s["pair_i_rad"]) for s in structures], dtype=np.int64)
- n_ang = np.array([len(s["pair_i_ang"]) for s in structures], dtype=np.int64)
- self.natoms = [int(s["natoms"]) for s in structures]
-
- # preprocess_structures already returns arrays with the right dtype
- # (int64 for indices, float32 for rij). Skip the defensive astype —
- # it was creating an extra copy of every per-frame array.
- at_cat = np.concatenate([s["atom_types"] for s in structures])
- pi_r_cat = np.concatenate([s["pair_i_rad"] for s in structures])
- pj_r_cat = np.concatenate([s["pair_j_rad"] for s in structures])
- rij_r_cat = np.concatenate([s["rij_rad"] for s in structures])
- pi_a_cat = np.concatenate([s["pair_i_ang"] for s in structures])
- pj_a_cat = np.concatenate([s["pair_j_ang"] for s in structures])
- rij_a_cat = np.concatenate([s["rij_ang"] for s in structures])
-
- at_all = torch.from_numpy(at_cat).to(device=device, non_blocking=True)
-
- pi_r_all = torch.from_numpy(pi_r_cat).to(device=device, non_blocking=True)
- pj_r_all = torch.from_numpy(pj_r_cat).to(device=device, non_blocking=True)
- rij_r_all = torch.from_numpy(rij_r_cat).to(device=device, dtype=dtype,
- non_blocking=True)
- pi_a_all = torch.from_numpy(pi_a_cat).to(device=device, non_blocking=True)
- pj_a_all = torch.from_numpy(pj_a_cat).to(device=device, non_blocking=True)
- rij_a_all = torch.from_numpy(rij_a_cat).to(device=device, dtype=dtype,
- non_blocking=True)
-
- self.energy = [float(s["energy"]) if "energy" in s else 0.0
- for s in structures]
- self.has_energy_flag = ["energy" in s for s in structures]
- self.has_forces_flag = ["forces" in s for s in structures]
- self.has_virial_flag = ["virial" in s for s in structures]
-
- f_parts = []
- for s in structures:
- if "forces" in s:
- f_parts.append(np.asarray(s["forces"]).reshape(-1, 3))
- else:
- f_parts.append(np.zeros((s["natoms"], 3), dtype=np.float32))
- f_cat = np.concatenate(f_parts).astype(np.float32 if dtype == torch.float32
- else np.float64, copy=False)
- f_all = torch.from_numpy(f_cat).to(device=device, dtype=dtype,
- non_blocking=True)
-
- v_parts = []
- for s in structures:
- if "virial" in s:
- v = np.asarray(s["virial"]).reshape(-1)
- if v.shape[0] == 6:
- v9 = np.array([v[0], v[3], v[5],
- v[3], v[1], v[4],
- v[5], v[4], v[2]])
- v_parts.append(v9)
- else:
- v_parts.append(v[:9])
- else:
- v_parts.append(np.zeros(9))
- v_cat = np.stack(v_parts).astype(np.float32 if dtype == torch.float32
- else np.float64, copy=False)
- v_all = torch.from_numpy(v_cat).to(device=device, dtype=dtype,
- non_blocking=True)
-
- # Per-frame cell volume (A**3) — needed for stress RMSE. Same order as
- # frames, so a batch slice follows the same indexing as .energy etc.
- vol_cat = np.asarray([s.get("volume", 0.0) for s in structures],
- dtype=np.float32 if dtype == torch.float32
- else np.float64)
- self.volumes = torch.from_numpy(vol_cat).to(device=device, dtype=dtype,
- non_blocking=True)
-
- if config is not None:
- # Build the cached basis in pair-chunks, writing into preallocated
- # buffers. The transient working set is then one chunk, not the
- # whole shard — this caps the construction-time GPU memory peak so
- # it stays close to the steady training footprint. Chebyshev and
- # angular bases are per-pair elementwise, so this is bit-identical
- # to computing them in one shot; only __init__ does more work, the
- # training step (which reads the split views below) is unchanged.
- rc_r = config["cutoff_radial"]
- rc_a = config["cutoff_angular"]
- bs_r = config["basis_size_radial"]
- bs_a = config["basis_size_angular"]
- l3 = config["l_max"][0]
- num_lm = sum(2 * ll + 1 for ll in range(1, l3 + 1)) if l3 >= 1 else 0
-
- P_r = rij_r_all.shape[0]
- fk_r_all = torch.empty(P_r, bs_r + 1, dtype=dtype, device=device)
- fkp_r_all = torch.empty(P_r, bs_r + 1, dtype=dtype, device=device)
- d12inv_r_all = torch.empty(P_r, dtype=dtype, device=device)
-
- P_a = rij_a_all.shape[0]
- fk_a_all = torch.empty(P_a, bs_a + 1, dtype=dtype, device=device)
- fkp_a_all = torch.empty(P_a, bs_a + 1, dtype=dtype, device=device)
- d12inv_a_all = torch.empty(P_a, dtype=dtype, device=device)
- blm_all = torch.empty(P_a, num_lm, dtype=dtype, device=device)
-
- # Chunk sized against memory still free *after* the buffers above.
- chunk = _basis_chunk_size(device, dtype, bs_a, num_lm, l3)
-
- for st in range(0, P_r, chunk):
- en = min(st + chunk, P_r)
- dr = torch.norm(rij_r_all[st:en], dim=-1)
- fk, fkp = ops.chebyshev_basis_and_deriv(dr, rc_r, bs_r)
- fk_r_all[st:en] = fk
- fkp_r_all[st:en] = fkp
- d12inv_r_all[st:en] = 1.0 / dr
-
- for st in range(0, P_a, chunk):
- en = min(st + chunk, P_a)
- rij = rij_a_all[st:en]
- da = torch.norm(rij, dim=-1)
- fk, fkp = ops.chebyshev_basis_and_deriv(da, rc_a, bs_a)
- fk_a_all[st:en] = fk
- fkp_a_all[st:en] = fkp
- dinv = 1.0 / da
- d12inv_a_all[st:en] = dinv
- if num_lm > 0:
- blm_all[st:en] = ops.angular_basis(
- rij[:, 0] * dinv, rij[:, 1] * dinv,
- rij[:, 2] * dinv, l3)
-
- nr_list = n_rad.tolist()
- na_list = n_ang.tolist()
- nat_list = [int(x) for x in self.natoms]
-
- self.atom_types = list(torch.split(at_all, nat_list))
- self.pi_rad = list(torch.split(pi_r_all, nr_list))
- self.pj_rad = list(torch.split(pj_r_all, nr_list))
- self.rij_rad = list(torch.split(rij_r_all, nr_list))
- self.pi_ang = list(torch.split(pi_a_all, na_list))
- self.pj_ang = list(torch.split(pj_a_all, na_list))
- self.rij_ang = list(torch.split(rij_a_all, na_list))
- self.forces = list(torch.split(f_all, nat_list))
- self.virial = list(torch.unbind(v_all, dim=0))
-
- if config is not None:
- self.fk_rad = list(torch.split(fk_r_all, nr_list))
- self.fkp_rad = list(torch.split(fkp_r_all, nr_list))
- self.d12inv_rad = list(torch.split(d12inv_r_all, nr_list))
- self.fk_ang = list(torch.split(fk_a_all, na_list))
- self.fkp_ang = list(torch.split(fkp_a_all, na_list))
- self.d12inv_ang = list(torch.split(d12inv_a_all, na_list))
- self.blm = list(torch.split(blm_all, na_list))
-
- self.n_energy = sum(self.has_energy_flag)
- self.n_forces = sum(self.has_forces_flag)
- self.n_virial = sum(self.has_virial_flag)
- self.has_forces = self.n_forces > 0
- self.has_virial = self.n_virial > 0
-
- def collate(self, indices: List[int]) -> Dict:
- """Fast GPU-side batch collation. No CPU->GPU transfer."""
- offsets = [0]
- for i in indices:
- offsets.append(offsets[-1] + self.natoms[i])
- N_total = offsets[-1]
- B = len(indices)
-
- at_list = [self.atom_types[i] for i in indices]
- atom_types = torch.cat(at_list)
-
- struct_idx = torch.cat([
- torch.full((self.natoms[i],), k, dtype=torch.long,
- device=self.device)
- for k, i in enumerate(indices)
- ])
-
- pi_r = torch.cat([self.pi_rad[i] + offsets[k]
- for k, i in enumerate(indices)])
- pj_r = torch.cat([self.pj_rad[i] + offsets[k]
- for k, i in enumerate(indices)])
- rij_r = torch.cat([self.rij_rad[i] for i in indices])
- pi_a = torch.cat([self.pi_ang[i] + offsets[k]
- for k, i in enumerate(indices)])
- pj_a = torch.cat([self.pj_ang[i] + offsets[k]
- for k, i in enumerate(indices)])
- rij_a = torch.cat([self.rij_ang[i] for i in indices])
-
- energy = torch.tensor([self.energy[i] for i in indices],
- dtype=self.dtype, device=self.device)
- natoms = torch.tensor([self.natoms[i] for i in indices],
- dtype=self.dtype, device=self.device)
-
- volumes = self.volumes[torch.as_tensor(indices, device=self.device,
- dtype=torch.long)]
-
- batch = {
- "N": N_total, "num_structures": B,
- "atom_types": atom_types, "struct_idx": struct_idx,
- "pair_i_rad": pi_r, "pair_j_rad": pj_r, "rij_rad": rij_r,
- "pair_i_ang": pi_a, "pair_j_ang": pj_a, "rij_ang": rij_a,
- "energy": energy, "natoms": natoms, "volumes": volumes,
- }
-
- batch["energy_mask"] = torch.tensor(
- [self.has_energy_flag[i] for i in indices],
- dtype=torch.bool, device=self.device)
-
- batch["forces"] = torch.cat([self.forces[i] for i in indices])
- force_flags = [self.has_forces_flag[i] for i in indices]
- batch["force_mask"] = torch.cat([
- torch.full((self.natoms[indices[k]],), force_flags[k],
- dtype=torch.bool, device=self.device)
- for k in range(B)
- ])
-
- batch["virial"] = torch.stack([self.virial[i] for i in indices])
- batch["virial_mask"] = torch.tensor(
- [self.has_virial_flag[i] for i in indices],
- dtype=torch.bool, device=self.device)
-
- if self.has_cached_basis:
- batch["fk_rad"] = torch.cat([self.fk_rad[i] for i in indices])
- batch["fkp_rad"] = torch.cat([self.fkp_rad[i] for i in indices])
- batch["d12inv_rad"] = torch.cat([self.d12inv_rad[i] for i in indices])
- batch["fk_ang"] = torch.cat([self.fk_ang[i] for i in indices])
- batch["fkp_ang"] = torch.cat([self.fkp_ang[i] for i in indices])
- batch["d12inv_ang"] = torch.cat([self.d12inv_ang[i] for i in indices])
- batch["blm"] = torch.cat([self.blm[i] for i in indices])
-
- return batch
-
-
class StreamDataStore:
- """Host-resident data store — GPU memory scales with batch size only.
-
- Same construction signature and same consumer interface as
- ``GPUDataStore`` (``collate`` + the metadata attributes), but every
- per-structure array stays in host memory. ``collate`` assembles the
- requested frames on the CPU, copies just that batch to the device
- (via a pinned staging copy so the H2D transfer is async), and computes
- the Chebyshev/angular basis for the batch on the fly.
-
- The recomputed basis is bit-identical to GPUDataStore's cached one:
- the basis is per-pair elementwise math, so evaluating it per batch
- reproduces exactly the values the cached path would have stored.
-
- Trade-off: each batch pays one small host->device transfer plus the
- basis recompute; in exchange the GPU never holds the dataset, so the
- memory footprint is set by the batch size, not the dataset size.
+ """The training data store: host-resident, GPU memory scales with batch
+ size only.
+
+ Every per-structure array stays in host memory. ``collate(indices)``
+ assembles the requested frames on the CPU, copies just that batch to the
+ device (via a pinned staging copy so the H2D transfer is async), and
+ computes the Chebyshev/angular basis for the batch on the fly. The
+ background-prefetch iteration in ``iter_collated`` plus the pinned async
+ copies hide the streaming work behind the device compute — benchmarks
+ across 1-16 element types showed speed parity with a fully preloaded
+ GPU store (which this class replaced) at ~10-15x less GPU memory, so
+ streaming is the only data path.
+
+ The per-batch basis is chunking-invariant per-pair elementwise math, so
+ every batch is bit-identical no matter how the dataset is batched.
"""
def __init__(self, structures: List[Dict], device: torch.device,
@@ -590,12 +311,12 @@ def __init__(self, structures: List[Dict], device: torch.device,
# Per-frame CPU views — predict_from_store reads
# ``data_store.forces[i].cpu()`` / ``.virial[i].cpu()`` (no-ops on
- # these CPU views), so the GPUDataStore interface is preserved.
+ # these CPU views).
self.forces = list(torch.split(self._f_all, self.natoms))
self.virial = list(torch.unbind(self._v_all, dim=0))
- # Tiny (n,) tensor; kept on device like GPUDataStore (collate indexes
- # it with a device tensor and predict calls .cpu() on it).
+ # Tiny (n,) tensor; kept on device (collate indexes it with a
+ # device tensor and predict calls .cpu() on it).
vol_cat = np.asarray([s.get("volume", 0.0) for s in structures],
dtype=np_dtype)
self.volumes = torch.from_numpy(vol_cat).to(device=device, dtype=dtype)
@@ -642,7 +363,7 @@ def _assemble_cpu(self, indices: List[int]) -> Dict:
off_t = torch.from_numpy(offsets[:-1])
# Pair indices are frame-local in storage; shift each frame's pairs
- # by its atom offset within the batch (same as GPUDataStore.collate).
+ # by its atom offset within the batch.
off_rep_r = torch.repeat_interleave(off_t, nr_t)
off_rep_a = torch.repeat_interleave(off_t, na_t)
@@ -682,11 +403,11 @@ def _stage(t):
"force_mask", "virial", "virial_mask")
def _basis_impl(self, rij_r, rij_a):
- """Per-batch basis — the values GPUDataStore would have cached.
+ """Per-batch Chebyshev/angular basis.
- Eager execution of this method is bit-identical to the cached path
- (per-pair elementwise math, chunking-invariant). ``compile_basis``
- may swap in a torch.compile'd version of this same method.
+ Per-pair elementwise math — chunking-invariant, so results do not
+ depend on how the dataset is batched. ``compile_basis`` may swap in
+ a torch.compile'd version of this same method.
"""
dr = torch.norm(rij_r, dim=-1)
fk_r, fkp_r = ops.chebyshev_basis_and_deriv(dr, self._rc_r, self._bs_r)
@@ -708,7 +429,7 @@ def compile_basis(self):
compiled: Inductor's fusion reassociates float ops, so the compiled
basis deviates from the eager one at the same ~1e-7 level the
compiled training compute already introduces. Never enabled in
- eager runs, which keeps the default bit-identical to GPUDataStore.
+ eager runs, which keeps eager training exactly chunking-invariant.
"""
self._basis_fn = torch.compile(self._basis_impl, dynamic=True)
@@ -731,8 +452,7 @@ def _finalize(self, staged: Dict) -> Dict:
def collate(self, indices: List[int]) -> Dict:
"""Assemble one batch on the CPU, ship it to the device, and compute
- the batch's basis there. Returns the same dict (same values) as
- ``GPUDataStore.collate``."""
+ the batch's basis there. Returns the collated batch dict."""
return self._finalize(self._assemble_cpu(indices))
@@ -1411,7 +1131,6 @@ def train_nep(
run_seed: int = None,
valid_file: str = None,
valid_ratio: float = None,
- stream_mode: bool = False,
):
"""Train a NEP model on a single device (GPU / CPU / MPS).
@@ -1488,16 +1207,6 @@ def train_nep(
run_seed, so it is reproducible for a given seed and is preserved
exactly on resume (the checkpoint's seed wins). Mutually exclusive
with valid_file.
- stream_mode : False (default) -> all structure data and the cached basis
- are pre-loaded to the device (fastest; GPU memory grows with dataset
- size). True -> the dataset stays in host memory and only the current
- batch is shipped to the device, with its basis computed on the fly —
- GPU memory then scales with ``batch`` instead of dataset size, at a
- modest per-epoch speed cost. Eager runs are numerically identical
- to the default (the streamed batches are bit-identical to the
- preloaded ones); with ``use_compile=True`` the per-batch basis is
- compiled too, adding only the same ~1e-7-level deviations the
- compiled compute already introduces.
"""
_clean_warning_format()
@@ -1644,7 +1353,7 @@ def _log(msg=""):
n_structs = len(frames)
# slim_types: detect which element types actually appear in the data and
- # narrow config before building neighbor lists / GPUDataStore / model.
+ # narrow config before building neighbor lists / data store / model.
# This makes the entire training run faster, not just the output file.
_slim_keep = None # None = no slimming; list = types to keep
if slim_types:
@@ -1672,16 +1381,12 @@ def _log(msg=""):
max_NN_rad, max_NN_ang = compute_max_neighbors(structures)
t0 = time.time()
- _StoreCls = StreamDataStore if stream_mode else GPUDataStore
- data_store = _StoreCls(structures, dev, dtype, config=config)
+ data_store = StreamDataStore(structures, dev, dtype, config=config)
del structures
if dev.type == "cuda":
torch.cuda.synchronize()
- if stream_mode:
- _log(f" stream_mode: dataset kept in host memory, batches "
- f"streamed to {dev} ({time.time() - t0:.1f}s)")
- else:
- _log(f" loaded to {dev} in {time.time() - t0:.1f}s (cached basis)")
+ _log(f" data store ready: dataset in host memory, batches streamed "
+ f"to {dev} ({time.time() - t0:.1f}s)")
_log(f" coverage: {data_store.n_energy} E / "
f"{data_store.n_forces} F / {data_store.n_virial} V")
@@ -1695,7 +1400,7 @@ def _log(msg=""):
vNN_rad, vNN_ang = compute_max_neighbors(structures_v)
max_NN_rad = max(max_NN_rad, vNN_rad)
max_NN_ang = max(max_NN_ang, vNN_ang)
- valid_store = _StoreCls(structures_v, dev, dtype, config=config)
+ valid_store = StreamDataStore(structures_v, dev, dtype, config=config)
del structures_v, valid_frames
_log(f" validation set ready in {time.time() - t0:.1f}s — "
f"coverage: {valid_store.n_energy} E / "
@@ -1798,8 +1503,13 @@ def _load_weights(target, path):
if not ok:
compile_msg = f" torch.compile: disabled — {msg}"
elif use_autograd_forces:
- compile_msg = (" torch.compile: skipped — incompatible with "
- "autograd double-backward forces")
+ # The nested create_graph=True double backward cannot be lowered
+ # directly, but the make_fx route (materialize the first-order
+ # gradient as ordinary FX ops, then compile) can — see
+ # torchnep/compiled_autograd.py.
+ compile_on = True
+ compile_msg = (" torch.compile: enabled (autograd forces via "
+ "make_fx-materialized gradient)")
else:
compile_on = True
compile_msg = " torch.compile: enabled (analytical compute method)"
@@ -1833,12 +1543,11 @@ def _load_weights(target, path):
"(recompute_q_scaler=True) — the loaded weights will "
"see rescaled descriptors and must re-adapt")
t0 = time.time()
- # In stream mode the q-scaler pass uses the training batch size so
- # its transient GPU footprint stays batch-bound (min/max over the
- # dataset is chunking-invariant, so the result is unchanged).
- qscaler_bs = batch_size if stream_mode else 1000
+ # The q-scaler pass uses the training batch size so its transient
+ # GPU footprint stays batch-bound (min/max over the dataset is
+ # chunking-invariant, so the result is unchanged).
q_min, q_max = compute_q_scaler(model, data_store, backend=backend,
- batch_size=qscaler_bs,
+ batch_size=batch_size,
gpumd_init=use_gpumd_qscaler)
model.set_q_scaler(q_min, q_max)
if dev.type == "cuda":
@@ -1855,11 +1564,21 @@ def _load_weights(target, path):
compute_props_cached = raw_model.compute_properties_cached
if compile_on:
_quiet_compile_logs()
- compute_props_cached = torch.compile(
- raw_model.compute_properties_cached, dynamic=True)
- if stream_mode:
+ if use_autograd_forces:
+ from .compiled_autograd import CompiledAutogradForce
+ compute_props = CompiledAutogradForce(raw_model).compute_properties
+ else:
+ # Compile the branch-free core only — ZBL and result assembly
+ # stay eager in the wrapper. Compiling the whole method instead
+ # leaves 3 graph breaks (per-type masks + ZBL autograd.grad) and
+ # ~2.7x the kernel launches.
+ import functools
+ _compiled_core = torch.compile(raw_model._cached_core,
+ dynamic=True)
+ compute_props_cached = functools.partial(
+ raw_model.compute_properties_cached, core_fn=_compiled_core)
# Fuse the per-batch basis kernels too — same numerical status
- # as the compiled compute (Inductor-level ~1e-7 deviations).
+ # as the compiled compute (~1e-7 Inductor deviations).
data_store.compile_basis()
if valid_store is not None:
valid_store.compile_basis()
diff --git a/torchnep/train_sharded.py b/torchnep/train_sharded.py
index b1fd424..1afc7af 100644
--- a/torchnep/train_sharded.py
+++ b/torchnep/train_sharded.py
@@ -98,8 +98,12 @@ def __init__(self, model: NEPModel, use_compile: bool = False):
# torch.compile's donated-buffer optimisation). See train.py for the
# single-device counterpart.
if use_compile and hasattr(torch, "compile"):
- self._compute_cached = torch.compile(
- model.compute_properties_cached, dynamic=True)
+ # Compile the branch-free core; ZBL + assembly stay eager in the
+ # wrapper (see NEPModel._cached_core).
+ import functools
+ self._compute_cached = functools.partial(
+ model.compute_properties_cached,
+ core_fn=torch.compile(model._cached_core, dynamic=True))
else:
self._compute_cached = model.compute_properties_cached
@@ -120,7 +124,7 @@ def forward(self, batch, use_autograd_forces: bool,
from .train import (
_BANNER, _AUTHOR,
- _backend_info, GPUDataStore, StreamDataStore, iter_collated,
+ _backend_info, StreamDataStore, iter_collated,
format_config_summary,
preprocess_structures,
_save_checkpoint, _load_checkpoint,
@@ -216,7 +220,6 @@ def train_nep_sharded(
run_seed: int = None,
valid_file: str = None,
valid_ratio: float = None,
- stream_mode: bool = False,
):
"""Data-sharded NEP training. Launch via torchrun (or any launcher that
sets RANK / LOCAL_RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT).
@@ -251,12 +254,11 @@ def train_nep_sharded(
training frames; the error sums are all-reduced, so every rank sees the
identical validation loss (schedulers stay in lock-step).
- ``stream_mode`` mirrors ``train_nep``: each rank keeps its shard in host
- memory and streams only the current batch to its GPU (basis computed on
- the fly, CPU assembly prefetched one batch ahead) — per-GPU memory then
- scales with ``batch`` instead of shard size. Combines with sharding:
- ranks stream independent shards, DDP collectives are untouched (same
- step count per rank, gradients all-reduced as usual).
+ Each rank keeps its shard in host memory and streams only the current
+ batch to its GPU (basis computed on the fly, CPU assembly prefetched
+ one batch ahead) — per-GPU memory scales with ``batch``, not shard
+ size. DDP collectives are untouched: same step count per rank,
+ gradients all-reduced as usual.
"""
_clean_warning_format()
@@ -540,20 +542,16 @@ def _compute_max_neighbors_local(structures):
max_NN_rad, max_NN_ang = int(nn_t[0].item()), int(nn_t[1].item())
t0 = time.time()
- _StoreCls = StreamDataStore if stream_mode else GPUDataStore
- data_store = _StoreCls(structures, dev, dtype, config=config)
+ data_store = StreamDataStore(structures, dev, dtype, config=config)
del structures
valid_store = None
if structures_v is not None:
- valid_store = _StoreCls(structures_v, dev, dtype, config=config)
+ valid_store = StreamDataStore(structures_v, dev, dtype, config=config)
del structures_v
if cuda_available:
torch.cuda.synchronize()
- if stream_mode:
- _log(f" stream_mode: shard kept in host memory, batches "
- f"streamed to {dev} ({time.time() - t0:.1f}s)")
- else:
- _log(f" loaded to {dev} in {time.time() - t0:.1f}s (cached basis)")
+ _log(f" data store ready: shard in host memory, batches streamed "
+ f"to {dev} ({time.time() - t0:.1f}s)")
# Aggregate data counts across all ranks for the banner
counts_t = torch.tensor(
@@ -696,13 +694,13 @@ def _load_weights(target, path):
_log(" q_scaler: RECOMPUTED from the new dataset "
"(recompute_q_scaler=True) — the loaded weights will "
"see rescaled descriptors and must re-adapt")
- # q_scaler: local shard -> all_reduce. In stream mode the pass uses
- # the training batch size so its transient GPU footprint stays
- # batch-bound (min/max is chunking-invariant — result unchanged).
+ # q_scaler: local shard -> all_reduce. The pass uses the training
+ # batch size so its transient GPU footprint stays batch-bound
+ # (min/max is chunking-invariant — result unchanged).
t_qs = time.time()
q_min, q_max = _compute_q_scaler_sharded(
model, data_store, backend=backend,
- batch_size=batch_size if stream_mode else 1000,
+ batch_size=batch_size,
gpumd_init=use_gpumd_qscaler)
model.set_q_scaler(q_min, q_max)
if cuda_available:
@@ -720,12 +718,11 @@ def _load_weights(target, path):
_log(compile_msg)
if compile_on:
_quiet_compile_logs()
- if stream_mode:
- # Fuse the per-batch basis kernels too — same numerical status
- # as the compiled compute (Inductor-level ~1e-7 deviations).
- data_store.compile_basis()
- if valid_store is not None:
- valid_store.compile_basis()
+ # Fuse the per-batch basis kernels too — same numerical status
+ # as the compiled compute (Inductor-level ~1e-7 deviations).
+ data_store.compile_basis()
+ if valid_store is not None:
+ valid_store.compile_basis()
shim = _NEPDDPShim(model, use_compile=compile_on)
# All per-type nets are always touched in compute_properties_cached (dummy
# pass for types absent in a given batch) so DDP sees every parameter in