Stream mode + Compiled autograd

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 |
 

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
 

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`.
 

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):

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 <http://www.gnu.org/licenses/>.
+
+"""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

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