From d545f6315108e2763c62cba325b1c06e490ae29a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 30 Jul 2026 17:09:17 +0800 Subject: [PATCH 1/9] feat(lmdb): batch frames of unequal atom count LMDB batches previously required every frame to have the same atom count, which can leave sparse size groups under-filled and give their frames a disproportionate optimizer weight. Add `batch_size: "mix:N"` so frames of different sizes can share one atom-budgeted batch. Use two layouts according to the model contract. Eligible graph models consume one concatenated node axis with per-frame `n_node` counts; other models retain rectangular batches whose shorter frames are padded with `atype = -1`. Keep native-spin models on the rectangular public path because their output translation is spin-specific. Exclude phantom rows from neighbor graphs and model evaluation, scatter per-atom outputs back at the public boundary, and make losses and validation weight only real atoms. Consolidate LMDB sampling and decoding around an explicit batch layout so serial and worker-process decoding preserve the same field shapes and frame order. Cover the ragged training path with the existing DPA1 graph lower, alongside padding, compaction, loss-reduction, sampler, and decoder regressions. --- deepmd/dpmodel/loss/dos.py | 4 +- deepmd/dpmodel/loss/ener.py | 93 +- deepmd/dpmodel/loss/reduction.py | 109 +- deepmd/dpmodel/loss/tensor.py | 2 +- deepmd/dpmodel/model/make_model.py | 40 +- deepmd/dpmodel/utils/__init__.py | 8 +- deepmd/dpmodel/utils/batch.py | 13 +- deepmd/dpmodel/utils/lmdb_data.py | 1076 ++++++++++++++--- .../dpmodel/utils/neighbor_graph/__init__.py | 4 + .../utils/neighbor_graph/ase_builder.py | 2 +- .../dpmodel/utils/neighbor_graph/from_ijs.py | 38 +- deepmd/dpmodel/utils/neighbor_graph/graph.py | 123 ++ deepmd/pt/loss/dens.py | 86 +- deepmd/pt/loss/dos.py | 4 +- deepmd/pt/loss/ener.py | 36 +- deepmd/pt/loss/tensor.py | 2 +- deepmd/pt/model/model/sezm_model.py | 56 +- .../pt/model/model/sezm_native_spin_model.py | 32 +- deepmd/pt/train/training.py | 34 +- deepmd/pt/utils/lmdb_dataset.py | 95 +- deepmd/pt/utils/nv_nlist.py | 15 + deepmd/pt_expt/model/ener_model.py | 87 +- deepmd/pt_expt/model/make_model.py | 140 ++- deepmd/pt_expt/train/training.py | 188 ++- deepmd/pt_expt/train/wrapper.py | 19 +- deepmd/pt_expt/utils/edge_schema.py | 60 +- deepmd/pt_expt/utils/graph_builder.py | 80 ++ deepmd/pt_expt/utils/lmdb_dataset.py | 80 +- deepmd/pt_expt/utils/nv_graph_builder.py | 86 +- deepmd/pt_expt/utils/vesin_graph_builder.py | 4 +- deepmd/pt_expt/utils/vesin_neighbor_list.py | 8 + deepmd/utils/argcheck.py | 4 +- deepmd/utils/data_system.py | 6 + doc/data/system.md | 4 + doc/train/training-advanced.md | 15 + source/tests/common/dpmodel/test_from_ijs.py | 4 +- .../tests/common/dpmodel/test_graph_ragged.py | 60 + source/tests/common/dpmodel/test_lmdb_data.py | 505 +++++++- .../tests/common/dpmodel/test_loss_padding.py | 180 ++- .../common/dpmodel/test_loss_reduction.py | 72 +- source/tests/consistent/test_lmdb_data.py | 30 +- source/tests/pt/model/test_sezm_model.py | 186 ++- source/tests/pt/test_lmdb_dataloader.py | 275 ++++- source/tests/pt/test_loss_padding.py | 90 +- source/tests/pt_expt/test_lmdb_training.py | 188 +++ .../pt_expt/utils/test_nv_matrix_decode.py | 159 ++- 46 files changed, 3692 insertions(+), 710 deletions(-) diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index c943ff2aa4..45229640cc 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -159,7 +159,7 @@ def call( ) diff3d = local_pred - local_label # [nf, natoms, numb_dos] if "mask" in model_dict: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=numb_dos). maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] l2_local_loss_dos = masked_atom_mean( xp.square(diff3d), maskf, self.numb_dos @@ -184,7 +184,7 @@ def call( ) diff3d = local_pred_cdf - local_label_cdf # [nf, natoms, numb_dos] if "mask" in model_dict: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=numb_dos). maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] l2_local_loss_cdf = masked_atom_mean( xp.square(diff3d), maskf, self.numb_dos diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 34f7858b83..46fc67af51 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -280,18 +280,25 @@ def call( atom_pref, ) - # Per-frame mask: recover real-atom count per frame when mask is provided. - # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghosts. - if "mask" in model_dict: + # Two things about a batch decide how its terms reduce, and the node + # axis states them differently. + # + # ``inv``, the reciprocal real atom count of each frame, is what the + # extensive frame-level terms (energy, virial) divide by. ``maskf`` + # marks the padded rows the per-atom terms must skip. A rectangular + # batch carries both in its mask: summing it gives the counts, and its + # zeros are the padding. A ragged batch pads nothing, so it states the + # counts alone and its per-atom terms reduce over their whole axis. + maskf = None + inv = None + if "n_node" in model_dict: + inv = 1.0 / xp.astype(model_dict["n_node"], energy.dtype) # [nf] + elif "mask" in model_dict: maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] - real_natoms = xp.sum(maskf, axis=-1) # [nf] - inv = xp.reshape(1.0 / real_natoms, (-1,)) # [nf] - _nf = maskf.shape[0] + inv = xp.reshape(1.0 / xp.sum(maskf, axis=-1), (-1,)) # [nf] _nloc = maskf.shape[1] - else: - # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, - # so leaving them unset here is safe (and avoids dead-store warnings). - maskf = None + if inv is not None: + _nf = inv.shape[0] if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as @@ -351,7 +358,7 @@ def call( if self.has_e: if self.loss_func == "mse": l2_ener_loss = xp.mean(xp.square(energy - energy_hat)) - if maskf is not None: + if inv is not None: # Idiom 2 (extensive): per-frame normalization by real-atom count. se = xp.square(energy - energy_hat) # [nf, k] per_frame = per_frame_component_mean(se) # [nf] @@ -383,7 +390,7 @@ def call( ) elif self.loss_func == "mae": l1_ener_loss = xp.mean(xp.abs(energy - energy_hat)) - if maskf is not None: + if inv is not None: abs_e = xp.abs(energy - energy_hat) # [nf, k] per_frame_ae = per_frame_component_mean(abs_e) # [nf] l1_ener_masked = xp.mean(per_frame_ae * inv) @@ -401,7 +408,7 @@ def call( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - if maskf is not None: + if inv is not None: per_frame_ae = per_frame_component_mean(xp.abs(energy - energy_hat)) mae_e = xp.mean(per_frame_ae * inv) else: @@ -415,12 +422,16 @@ def call( if maskf is not None: # Idiom 1 (per-atom masked mean, ncomp=3). diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3] - maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] # Masked MSE computed for rmse_f display regardless of use_huber. l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3) if not self.use_huber: loss += pref_f * l2_force_masked else: + # ``f_use_norm`` selects the residual an atom + # contributes: three independent components, or the + # single L2 norm of its force-error vector. That choice + # sets the label count per atom, which is exactly the + # ``ncomp`` the pooled reduction divides by. if not self.f_use_norm: abs_e = xp.abs(diff_f_3d) quad = 0.5 * xp.square(diff_f_3d) @@ -429,8 +440,8 @@ def call( ) huber_elem = xp.where( abs_e <= self._huber_delta_force, quad, lin - ) - huber_masked = huber_elem * maskf_col + ) # [nf, nloc, 3] + huber_ncomp = 3 else: diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) norm_2d = xp.reshape( @@ -444,18 +455,16 @@ def call( lin_n = self._huber_delta_force * ( abs_n - 0.5 * self._huber_delta_force ) - huber_n = xp.where( - abs_n <= self._huber_delta_force, quad_n, lin_n + huber_elem = xp.reshape( + xp.where( + abs_n <= self._huber_delta_force, quad_n, lin_n + ), + (_nf, _nloc, 1), ) - huber_masked = xp.reshape(huber_n * maskf, (_nf, _nloc, 1)) - per_frame_sum = xp.sum( - xp.reshape(huber_masked, (_nf, -1)), axis=-1 + huber_ncomp = 1 + l_huber_masked = masked_atom_mean( + huber_elem, maskf, huber_ncomp ) - if not self.f_use_norm: - per_frame_dof = xp.sum(maskf, axis=-1) * 3 - else: - per_frame_dof = xp.sum(maskf, axis=-1) - l_huber_masked = xp.mean(per_frame_sum / per_frame_dof) loss += pref_f * l_huber_masked more_loss["rmse_f"] = self.display_if_exist( xp.sqrt(l2_force_masked), find_force @@ -495,10 +504,10 @@ def call( xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), (_nf, _nloc), ) - masked_norm = norm_2d * maskf - per_frame_sum = xp.sum(masked_norm, axis=-1) - per_frame_dof = xp.sum(maskf, axis=-1) - l1_force_masked = xp.mean(per_frame_sum / per_frame_dof) + # One L2 norm per atom, hence one label per atom. + l1_force_masked = masked_atom_mean( + xp.reshape(norm_2d, (_nf, _nloc, 1)), maskf, 1 + ) loss += pref_f * l1_force_masked more_loss["mae_f"] = self.display_if_exist( l1_force_masked, find_force @@ -533,7 +542,7 @@ def call( l2_virial_loss = xp.mean( xp.square(virial_hat_reshape - virial_reshape), ) - if maskf is not None: + if inv is not None: # Idiom 2 (extensive, k=9): per-frame normalization. v2d = xp.reshape(virial, (_nf, 9)) v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) @@ -567,7 +576,7 @@ def call( ) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) - if maskf is not None: + if inv is not None: v2d = xp.reshape(virial, (_nf, 9)) v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) per_frame_v = per_frame_component_mean( @@ -588,7 +597,7 @@ def call( f"Loss type {self.loss_func} is not implemented for virial loss." ) if mae: - if maskf is not None: + if inv is not None: v2d = xp.reshape(virial, (_nf, 9)) v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) per_frame_v = per_frame_component_mean(xp.abs(v_hat_2d - v2d)) @@ -609,7 +618,6 @@ def call( # Idiom 1 (per-atom masked mean, ncomp=1). ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) - per_frame_dof = xp.sum(maskf, axis=-1) # [nf] l2_ae_masked = masked_atom_mean( xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 ) @@ -626,9 +634,9 @@ def call( huber_ae = xp.where( abs_ae <= self._huber_delta_energy, quad_ae, lin_ae ) - huber_ae_masked = huber_ae * maskf - per_frame_sum_h = xp.sum(huber_ae_masked, axis=-1) - l_huber_ae_masked = xp.mean(per_frame_sum_h / per_frame_dof) + l_huber_ae_masked = masked_atom_mean( + huber_ae[:, :, None], maskf, 1 + ) loss += pref_ae * l_huber_ae_masked more_loss["rmse_ae"] = self.display_if_exist( xp.sqrt(l2_ae_masked), find_atom_ener @@ -714,6 +722,17 @@ def call( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." ) if self.has_gf: + if maskf is None and inv is not None: + # ``natoms`` below is one number for the whole batch, which a + # padded batch can honour and a concatenated one cannot: its + # frames differ in atom count, so ``drdq``, stored per frame + # against a common atom axis, has no shape to take. + raise NotImplementedError( + "the generalized force loss requires every frame of a " + "batch to hold the same number of atoms; a batch whose " + "frames are concatenated cannot provide the common atom " + "axis its ``drdq`` label is stored against" + ) find_drdq = label_dict["find_drdq"] drdq = label_dict["drdq"] pref_gf = find_drdq * ( diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py index 2568913538..573b00684a 100644 --- a/deepmd/dpmodel/loss/reduction.py +++ b/deepmd/dpmodel/loss/reduction.py @@ -1,19 +1,60 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Shared masked per-frame reduction idioms for the loss modules. - -These helpers factor out the three per-frame reduction patterns that the -mixed_type padding mask (PR #5738) introduced into every loss term (issue -#5768). They are written with ``array_api_compat`` so both the dpmodel -(numpy/jax/...) loss backend and the PyTorch loss backend can call them: the -PyTorch backend passes torch tensors and ``array_api_compat`` dispatches to the -torch namespace, preserving autograd and producing bit-identical results to the -previous hand-inlined torch code. - -Each helper implements ONLY the masked branch. Callers keep the original -non-masked expression in the ``else`` branch verbatim, so the "bit-identical -for non-mixed batches" guarantee from PR #5738 is preserved (defaulting the -mask to all-ones would change the reduction order at the ULP level and is -deliberately NOT done here). +"""Shared masked reduction idioms for the loss modules. + +These helpers factor out the reduction patterns every loss term needs once a +per-atom mask marks which rows of a batch are real. They are written with +``array_api_compat`` so both the dpmodel (numpy/jax/...) loss backend and the +PyTorch loss backend can call them: the PyTorch backend passes torch tensors +and ``array_api_compat`` dispatches to the torch namespace, preserving +autograd. + +Reduction convention +-------------------- +A batch may hold frames of unequal atom count, padded to a common width. Each +term must therefore decide how one frame's aggregate error weighs against +another's, and the answer differs by term because frames do not carry equally +many labels: + +- **Per-atom terms** (force, atomic energy, atomic prefactor force, dos, + tensor) carry a number of labels proportional to the frame's atom count. + :func:`masked_atom_mean` pools them: it divides the summed contribution of + the whole batch by the batch's total label count, so every real label counts + once and a frame's weight is proportional to its atom count. +- **Frame-level terms** (energy, virial) carry a fixed number of labels per + frame. Pooling and averaging over frames coincide there, so + :func:`per_frame_component_mean` reduces per frame and leaves the frame axis + to the caller, which applies the extensive ``1 / natoms`` weighting. + +Pooling is what keeps a frame's weight independent of the company it keeps. +The alternative -- averaging each frame's own per-label mean -- gives every +frame the same weight whatever its size, which makes a label in a small frame +count for more than one in a large frame, by the ratio of their atom counts. + +Writing ``S_f`` for the summed contribution of frame ``f``, ``k`` for the +number of frames and ``d`` for the labels each of them carries, the two +coincide whenever that count is common to the batch: + + sum_f(S_f) / (k * d) == (1 / k) * sum_f(S_f / d) + +so a batch of uniform atom count needs no special case anywhere in this +module, and the choice between the two is unobservable there. They part +company only where a batch holds frames of differing real atom count, which +arises in exactly two places: + +- ``mix:N`` LMDB batching, which packs frames of differing atom count by + construction. +- ``mixed_type`` npy data whose ``real_atom_types.npy`` spends a different + number of ``-1`` rows on different frames of one system. The format permits + this and the documentation describes it as the way to merge frames of + unequal atom count, but a system written by dpdata pads every frame equally + and is therefore unaffected. + +The TensorFlow backend reaches neither case: it drops ``real_natoms_vec`` +before the feed dict and normalizes by the padded width throughout. + +Each helper implements ONLY the masked branch. The unmasked branch of each +caller is a plain mean over the whole batch, which pools by construction, so +both branches express the same convention. """ from typing import ( @@ -28,7 +69,13 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: - """Idiom 1: per-atom masked mean over ``ncomp`` components, averaged over frames. + """Idiom 1: mean of a per-atom contribution over the batch's real labels. + + The contribution of every real atom is pooled across frames before the + division, so the reduction is a mean over labels rather than a mean over + frames of per-frame means. See the module docstring for why the per-atom + terms weigh frames by their label count, and for the identity that makes + this reduce to the per-frame mean on a uniform-atom-count batch. Parameters ---------- @@ -45,26 +92,20 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: Returns ------- Array - ``mean_over_frames( sum(elem * mask) / (real_natoms * ncomp) )``. - An all-padding frame (zero real atoms) contributes a neutral ``0`` - instead of ``0/0 = NaN``. + ``sum(elem * mask) / (ncomp * sum(mask))`` over the whole batch. + A batch holding no real atom contributes a neutral ``0`` instead of + ``0/0 = NaN``. """ xp = array_api_compat.array_namespace(elem, maskf) - nf = elem.shape[0] - masked = elem * maskf[:, :, None] - per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1) - per_frame_dof = xp.sum(maskf, axis=-1) * ncomp - # An all-padding frame has zero real atoms, so ``per_frame_dof`` is 0 and - # the ratio would be 0/0 = NaN -- poisoning the frame mean and, under - # autograd, its gradient. Divide by a safe denominator and map those frames - # to a neutral per-frame value of 0. Frames with real atoms are untouched, - # preserving the bit-identical guarantee. - has_dof = per_frame_dof > 0 - safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof)) - per_frame = xp.where( - has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum) - ) - return xp.mean(per_frame) + total = xp.sum(elem * maskf[:, :, None]) + total_dof = xp.sum(maskf) * ncomp + # A batch of nothing but padding has no label to average over, and the + # ratio would be 0/0 = NaN -- poisoning the whole batch loss and, under + # autograd, its gradient. The division still runs on a safe denominator so + # that the discarded branch stays differentiable. + has_dof = total_dof > 0 + safe_dof = xp.where(has_dof, total_dof, xp.ones_like(total_dof)) + return xp.where(has_dof, total / safe_dof, xp.zeros_like(total)) def masked_pair_mean(elem: Array, maskf: Array, ncomp: int) -> Array: diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index e054ab3cc6..11a8ce3987 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -124,7 +124,7 @@ def call( diff = xp.reshape(local_pred - local_label, (-1, self.tensor_size)) diff = diff * atomic_weight if "mask" in model_dict: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=tensor_size). maskf = xp.astype(model_dict["mask"], diff.dtype) # [nf, natoms] diff3d = xp.reshape( diff, (local_pred.shape[0], natoms, self.tensor_size) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 44967bba74..10d8f01e87 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -48,6 +48,8 @@ NeighborGraph, build_neighbor_graph, build_neighbor_graph_ase, + compact_nodes, + expand_node_values, ) from deepmd.utils.path import ( DPPath, @@ -543,12 +545,21 @@ def _call_common_graph( ) xp = array_api_compat.array_namespace(atype) nf, nloc = atype.shape[:2] + n_padded = nf * nloc + atype_flat = xp.reshape(atype, (n_padded,)) + # A batch of unequal atom counts arrives padded to a common width + # with phantom atoms (atype < 0). The builders leave them out of + # every edge, so dropping them from the node axis costs nothing and + # spares the network from evaluating them. On a batch of uniform + # atom count the mask is all true and this is a renumbering by the + # identity. + ng, node_index = compact_nodes(ng, atype_flat >= 0) # OUTPUT-AGNOSTIC standard model dict (````, ``_redu``, # derivative name-holders ``None``, plus int ``mask``), like the # dense ``call_common``. ``call_lower_graph`` masks virtual atoms # (atype<0) and sets the real int mask. model_predict = self.call_lower_graph( - atype=xp.reshape(atype, (nf * nloc,)), + atype=xp.take(atype_flat, node_index, axis=0), n_node=ng.n_node, edge_index=ng.edge_index, edge_vec=ng.edge_vec, @@ -556,25 +567,38 @@ def _call_common_graph( fparam=fp, # graph-lower ABI: aparam is FLAT on the node axis, (N, nda). aparam=( - xp.reshape(ap, (nf * nloc, ap.shape[-1])) + xp.take( + xp.reshape(ap, (n_padded, ap.shape[-1])), node_index, axis=0 + ) if ap is not None else None ), - spin=(xp.reshape(spin, (nf * nloc, 3)) if spin is not None else None), + spin=( + xp.take(xp.reshape(spin, (n_padded, 3)), node_index, axis=0) + if spin is not None + else None + ), charge_spin=charge_spin, ) - # Public ABI is rectangular (nf, nloc, *); the lower is flat - # (N=nf*nloc, *). Unravel per-atom keys here at the boundary. - # public call_common always passes rectangular (nf,nloc) coord/atype (N == nf*nloc), so this unravel always applies; ragged graphs reach call_lower_graph/forward_common_lower_graph directly (no unravel) and stay flat (N,*). + # Public ABI is rectangular (nf, nloc, *); the lower is flat over + # the real atoms. Scatter per-atom keys back onto the padded width + # here at the boundary, where a phantom slot reads zero, which is + # what a masked-out atom contributed there before. + # Only the rectangular entry reaches this scatter; the ragged + # one keeps the flat axis its caller handed over. + n_real = node_index.shape[0] for k in list(model_predict.keys()): v = model_predict[k] # per-frame reduced keys (..._redu) keep their (nf, *) shape; only node-level (N,*) keys unravel — guards the nloc==1 case where N == nf. if ( v is not None and not k.endswith("_redu") - and v.shape[:1] == (nf * nloc,) + and v.shape[:1] == (n_real,) ): - model_predict[k] = xp.reshape(v, (nf, nloc, *v.shape[1:])) + model_predict[k] = xp.reshape( + expand_node_values(v, node_index, n_padded), + (nf, nloc, *v.shape[1:]), + ) return model_predict def call_common_lower( diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 3e439f173e..1f22690d1a 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -10,11 +10,11 @@ PairExcludeMask, ) from .lmdb_data import ( - DistributedSameNlocBatchSampler, + DistributedLmdbBatchSampler, + LmdbBatchSampler, LmdbDataReader, LmdbTestData, LmdbTestDataNlocView, - SameNlocBatchSampler, is_lmdb, make_neighbor_stat_data, ) @@ -79,11 +79,12 @@ __all__ = [ "AtomExcludeMask", "DefaultNeighborList", - "DistributedSameNlocBatchSampler", + "DistributedLmdbBatchSampler", "EmbeddingNet", "EnvMat", "FittingNet", "GraphLayout", + "LmdbBatchSampler", "LmdbDataReader", "LmdbTestData", "LmdbTestDataNlocView", @@ -93,7 +94,6 @@ "NeighborList", "NetworkCollection", "PairExcludeMask", - "SameNlocBatchSampler", "aggregate", "apply_pair_exclusion_nlist", "build_multiple_neighbor_list", diff --git a/deepmd/dpmodel/utils/batch.py b/deepmd/dpmodel/utils/batch.py index 2cbf8a72ff..77416d8dd6 100644 --- a/deepmd/dpmodel/utils/batch.py +++ b/deepmd/dpmodel/utils/batch.py @@ -11,7 +11,18 @@ _DROP_KEYS = {"default_mesh", "sid", "fid"} # Keys that belong to model input (everything else is label). -_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam", "charge_spin"} +# ``n_node`` is an input rather than a label: it states how a ragged batch's +# flat node axis divides into frames, which the model needs to read it at all. +_INPUT_KEYS = { + "coord", + "atype", + "spin", + "box", + "fparam", + "aparam", + "charge_spin", + "n_node", +} def normalize_batch(batch: dict[str, Any]) -> dict[str, Any]: diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 5cb22ed304..90cb0a5ae1 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -5,6 +5,7 @@ Backend-specific wrappers (PyTorch Dataset, JAX, etc.) import from here. """ +import dataclasses import logging import math import multiprocessing @@ -82,6 +83,38 @@ _LMDB_METADATA_KEYS = frozenset({"atom_numbs", "atom_names", "orig"}) _OPTIONAL_MODEL_INPUT_KEYS = frozenset({"fparam", "aparam", "spin", "charge_spin"}) +# Atom type written into the padded slots of a mixed-nloc batch. A phantom +# atom occupies a tensor slot but no physical site: the neighbor list gives it +# no neighbors, the atomic model zeroes its output, and the loss masks it out. +PHANTOM_ATOM_TYPE = -1 + +# Fields whose leading axis is the atom axis, and which a mixed-nloc batch must +# therefore pad to the batch-wide atom count. Membership cannot be inferred +# from array shapes: a frame with ``nloc == 9`` makes ``virial`` (shape ``(9,)``) +# indistinguishable from a per-atom field, and ``nloc == 2`` does the same for +# a two-component ``fparam``. The registered data requirements carry the +# authoritative ``atomic`` flag; these two sets cover the fields that exist +# without one. +_STRUCTURAL_PER_ATOM_KEYS = frozenset({"coord", "atype"}) +_OPTIONAL_PER_ATOM_KEYS = frozenset({"aparam", "spin"}) + +# Frame-level fields that a decoded frame may carry without a registered +# requirement. They anchor the shape-based fallback below, which classifies any +# remaining unrecognized field by comparing its leading axis to the frame's +# atom count. +_FRAME_LEVEL_KEYS = frozenset( + { + "box", + "energy", + "virial", + "fparam", + "charge_spin", + "natoms", + "real_natoms_vec", + "min_pair_dist", + } +) + # Process-level cache: python-lmdb does not allow opening the same path twice # in one process. We ref-count so the Environment is closed (and freed from # the cache) once every reader that shares it is garbage-collected. @@ -276,6 +309,62 @@ def _resolve_frame_dtype(config: LmdbDecodeConfig, key: str) -> np.dtype: return np.dtype(GLOBAL_NP_FLOAT_PRECISION) +def _requirement_is_atomic(requirement: Any) -> bool: + """Whether a data requirement describes a per-atom quantity.""" + if isinstance(requirement, dict): + return bool(requirement.get("atomic", False)) + return bool(getattr(requirement, "atomic", False)) + + +def resolve_per_atom_keys( + frame: dict[str, Any], + config: LmdbDecodeConfig, +) -> frozenset[str]: + """Return the fields of one frame whose leading axis is the atom axis. + + Classification is authoritative wherever possible: coordinates and atom + types are per-atom by construction, and every registered data requirement + declares whether it is ``atomic``. Only fields the loader has never been + told about fall back to comparing their leading axis against the atom + count, and the frame-level fields that DeePMD itself produces are excluded + from that fallback so a coincidental shape match cannot misclassify them. + + Parameters + ---------- + frame : dict[str, Any] + One decoded frame in DeePMD data-system convention. ``coord`` anchors + the atom axis: it is the one field present in every frame whose + leading axis is the atom count. + config : LmdbDecodeConfig + Decoder state holding the registered data requirements. + + Returns + ------- + frozenset[str] + Names of the fields to pad along their leading axis. + """ + nloc = frame["coord"].shape[0] + keys = set(_STRUCTURAL_PER_ATOM_KEYS) + keys |= { + key + for key, requirement in config.data_requirements.items() + if _requirement_is_atomic(requirement) + } + keys |= _OPTIONAL_PER_ATOM_KEYS + for key, value in frame.items(): + if ( + key in keys + or key in _FRAME_LEVEL_KEYS + or key in config.data_requirements + or key.startswith("find_") + or key == "fid" + ): + continue + if isinstance(value, np.ndarray) and value.ndim >= 1 and value.shape[0] == nloc: + keys.add(key) + return frozenset(keys & frame.keys()) + + def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray: """Build ``[nloc, nloc, count(type_0), ...]`` for one frame. @@ -460,11 +549,147 @@ def decode_lmdb_frame( return frame +def _pad_fill_value(key: str) -> int: + """Return the value written into the padded tail of one per-atom field. + + Atom types use the phantom sentinel so that downstream code recognizes the + slot as unoccupied; every other per-atom field is zeroed, which keeps + padded rows neutral in the sums and masked reductions that consume them. + """ + return PHANTOM_ATOM_TYPE if key == "atype" else 0 + + +def per_atom_strides( + frame: dict[str, Any], + per_atom_keys: frozenset[str], +) -> dict[str, int]: + """Return how many leading-axis entries each per-atom field spends per atom. + + Most per-atom fields carry one row per atom, so their leading axis is the + atom count itself. A data requirement declared with ``repeat != 1`` is + instead stored flat and atom-major, giving a leading axis of + ``nloc * repeat``. Padding widens that axis by whole atoms either way, so + the factor is all the padding logic needs to tell the two layouts apart. + + Parameters + ---------- + frame : dict[str, Any] + One decoded frame; ``coord`` anchors the atom count. + per_atom_keys : frozenset[str] + Fields to measure, as resolved by :func:`resolve_per_atom_keys`. + + Returns + ------- + dict[str, int] + Leading-axis entries per atom, keyed by field name. + + Raises + ------ + ValueError + If a field's leading axis is not a whole multiple of the atom count. + """ + nloc = frame["coord"].shape[0] + strides: dict[str, int] = {} + for key in per_atom_keys: + length = np.asarray(frame[key]).shape[0] + if nloc == 0 or length % nloc: + raise ValueError( + f"LMDB field {key!r} has a leading axis of {length}, which is " + f"not a whole number of entries per atom in a {nloc}-atom frame" + ) + strides[key] = length // nloc + return strides + + +@dataclass(frozen=True) +class BatchLayout: + """Where each frame's per-atom rows sit on a decoded batch's leading axis. + + Two layouts serve the two shapes a model's node axis can take. + + The **rectangular** layout gives every frame a row of the batch-wide atom + count and pads the tail of the shorter ones with phantom atoms, which is + what a model reading an ``(nf, nloc, ...)`` node axis requires. The + **ragged** layout concatenates the frames instead, so nothing is padded and + the leading axis is the batch's real atom count; a model reading a flat + node axis consumes that directly, paired with ``n_node``. Frame-level + fields are stacked on the frame axis either way. + + Attributes + ---------- + n_node : numpy.ndarray + Real atom count of each frame of this decode, with shape ``(nf,)``. + strides : dict[str, int] + Leading-axis entries per atom of each per-atom field, as resolved by + :func:`per_atom_strides`. + ragged : bool + Whether frames are concatenated rather than padded to a common width. + width : int + Atoms each frame occupies under the rectangular layout. It stays the + width of the whole batch even where ``n_node`` covers one chunk of it, + since chunks that padded to their own widths would not concatenate. + """ + + n_node: np.ndarray + strides: dict[str, int] + ragged: bool + width: int + + def __post_init__(self) -> None: + # Frame offsets on the ragged axis, in atoms. Held rather than summed + # per lookup, since every field of every frame asks for one. + object.__setattr__( + self, "_offset", np.concatenate([[0], np.cumsum(self.n_node)]) + ) + + @classmethod + def over( + cls, n_node: np.ndarray, strides: dict[str, int], *, ragged: bool + ) -> "BatchLayout": + """Return the layout of a batch holding the given per-frame counts.""" + return cls( + n_node=n_node, + strides=strides, + ragged=ragged, + width=int(n_node.max()) if n_node.size else 0, + ) + + def chunk(self, start: int, stop: int) -> "BatchLayout": + """Return the layout of a contiguous run of this batch's frames.""" + return dataclasses.replace(self, n_node=self.n_node[start:stop]) + + def field_length(self, key: str) -> int: + """Return the leading-axis length one per-atom field is allocated.""" + stride = self.strides[key] + if self.ragged: + return int(self.n_node.sum()) * stride + return self.width * stride + + def frame_index(self, row: int, key: str) -> Any: + """Return the index selecting one frame's rows of a per-atom field. + + Rectangular batches carry a frame axis, so the rows of frame ``row`` + are the head of its own row; ragged batches carry none, and the rows + are a run at the frame's offset. + """ + stride = self.strides[key] + length = int(self.n_node[row]) * stride + if not self.ragged: + return (row, slice(0, length)) + start = int(self._offset[row]) * stride + return slice(start, start + length) + + def _allocate_lmdb_batch( frame: dict[str, Any], batch_size: int, + layout: BatchLayout, ) -> dict[str, Any]: - """Allocate a contiguous NumPy batch from the first decoded frame.""" + """Allocate a contiguous NumPy batch from the first decoded frame. + + Per-atom fields are allocated at the length :class:`BatchLayout` gives + them: one padded row per frame, or one concatenated run over the batch. + """ batch: dict[str, Any] = {} for key, value in frame.items(): if key.startswith("find_"): @@ -476,6 +701,20 @@ def _allocate_lmdb_batch( continue elif value is None: batch[key] = None + elif key in layout.strides: + array = np.asarray(value) + head = ( + (layout.field_length(key),) + if layout.ragged + else (batch_size, layout.field_length(key)) + ) + destination = np.full( + (*head, *array.shape[1:]), + _pad_fill_value(key), + dtype=array.dtype, + ) + destination[layout.frame_index(0, key)] = array + batch[key] = destination else: array = np.asarray(value) destination = np.empty((batch_size, *array.shape), dtype=array.dtype) @@ -484,17 +723,72 @@ def _allocate_lmdb_batch( return batch +def _promote_batch_field( + batch: dict[str, Any], + field: str, + layout: BatchLayout, + frames_written: int, + dtype: np.dtype, +) -> np.ndarray: + """Widen one batch field's dtype in place, preserving what was written. + + The replacement is prefilled with the field's padding value rather than + zeroed, so that entries not yet written keep the marker the allocation gave + them; for ``atype`` under the rectangular layout that marker is what + identifies a phantom atom. + """ + destination = batch[field] + promoted = np.full(destination.shape, _pad_fill_value(field), dtype=dtype) + # The leading axis counts frames, except for a per-atom field of a ragged + # batch, where it counts the atom rows those frames have filled. + if field in layout.strides and layout.ragged: + written = int(layout._offset[frames_written]) * layout.strides[field] + else: + written = frames_written + promoted[:written] = destination[:written] + batch[field] = promoted + return promoted + + def decode_lmdb_batch( transaction: lmdb.Transaction, original_keys: Sequence[int], frame_format: str, config: LmdbDecodeConfig, + layout: BatchLayout | None = None, ) -> dict[str, Any]: """Decode LMDB records directly into preallocated contiguous arrays. The function keeps at most one temporary frame alive. It avoids the decode-copy, dtype-copy, Python frame-list, and final ``numpy.stack`` sequence used by generic collation. + + Parameters + ---------- + transaction : lmdb.Transaction + Open read transaction on the LMDB environment. + original_keys : Sequence[int] + Integer LMDB frame keys in batch order. + frame_format : str + Format specification for integer LMDB frame keys. + config : LmdbDecodeConfig + Decoder state independent of the LMDB environment. + layout : BatchLayout, optional + Where each frame's per-atom rows belong. Defaults to a rectangular + layout at the atom count of the first frame, which leaves a batch of + uniform atom count untouched. + + Returns + ------- + dict[str, Any] + One collated batch of contiguous NumPy arrays. A ragged layout adds + ``n_node``, the per-frame atom count its flat axis is read with. + + Notes + ----- + The layout fixes the shape of every field of the result, so a chunked + decode must pass each chunk the layout of its own frames, cut from the + batch-wide one; :meth:`LmdbDataReader.batch_layout` resolves that once. """ if not original_keys: raise ValueError("decode_lmdb_batch requires at least one frame key") @@ -513,15 +807,27 @@ def decode_lmdb_batch( config, copy_arrays=False, ) + frame_nloc = frame["coord"].shape[0] + if layout is None: + layout = BatchLayout.over( + np.full(batch_size, frame_nloc, dtype=np.int64), + per_atom_strides(frame, resolve_per_atom_keys(frame, config)), + ragged=False, + ) + if int(layout.n_node[row]) != frame_nloc: + raise ValueError( + f"the batch layout gives frame {original_key} " + f"{int(layout.n_node[row])} atoms, but it holds {frame_nloc}" + ) if batch is None: - batch = _allocate_lmdb_batch(frame, batch_size) + batch = _allocate_lmdb_batch(frame, batch_size, layout) expected_fields = frozenset(frame) continue frame_fields = frozenset(frame) if frame_fields != expected_fields: raise ValueError( - "LMDB frames in one same-nloc batch expose inconsistent fields: " + "LMDB frames in one batch expose inconsistent fields: " f"frame {original_keys[0]} has {sorted(expected_fields)}, while " f"frame {original_key} has {sorted(frame_fields)}" ) @@ -537,23 +843,45 @@ def decode_lmdb_batch( continue if field == "fid": batch[field][row] = value - else: - destination = batch[field] - array = np.asarray(value) - if destination.shape[1:] != array.shape: + continue + destination = batch[field] + array = np.asarray(value) + stride = layout.strides.get(field) + # A per-atom field may differ in its leading axis, which the layout + # absorbs as long as the axis stays a whole number of atoms; every + # remaining axis must match exactly. + if stride is not None: + lead_axes = 1 if layout.ragged else 2 + expected_tail: tuple[int, ...] = destination.shape[lead_axes:] + actual_tail = array.shape[1:] + if array.shape[0] != frame_nloc * stride: raise ValueError( - f"LMDB field {field!r} changes shape within one batch: " - f"expected {destination.shape[1:]}, got {array.shape} " - f"for frame {original_key}" + f"LMDB field {field!r} spends {array.shape[0]} leading " + f"entries on {frame_nloc} atoms in frame {original_key}, " + f"against {stride} per atom in frame {original_keys[0]}" ) - result_dtype = np.result_type(destination.dtype, array.dtype) - if result_dtype != destination.dtype: - promoted = np.empty(destination.shape, dtype=result_dtype) - promoted[:row] = destination[:row] - batch[field] = destination = promoted + else: + expected_tail = destination.shape[1:] + actual_tail = array.shape + if expected_tail != actual_tail: + raise ValueError( + f"LMDB field {field!r} changes shape within one batch: " + f"expected {expected_tail}, got {actual_tail} " + f"for frame {original_key}" + ) + result_dtype = np.result_type(destination.dtype, array.dtype) + if result_dtype != destination.dtype: + destination = _promote_batch_field( + batch, field, layout, row, result_dtype + ) + if stride is not None: + destination[layout.frame_index(row, field)] = array + else: destination[row] = array - assert batch is not None + assert batch is not None and layout is not None + if layout.ragged: + batch["n_node"] = layout.n_node batch["sid"] = np.asarray([0], dtype=np.int64) return batch @@ -569,8 +897,14 @@ def _decode_lmdb_worker_chunk( frame_format: str, config: LmdbDecodeConfig, original_keys: list[int], + layout: BatchLayout, ) -> dict[str, Any]: - """Decode one chunk using process-local LMDB state.""" + """Decode one chunk using process-local LMDB state. + + ``layout`` is this chunk's slice of the batch layout decided by the parent + process. Deriving it there rather than per chunk is what lets + :func:`_merge_lmdb_chunks` concatenate the results. + """ reader = _WORKER_LMDB_READERS.get(lmdb_path) if reader is None: environment = lmdb.open( @@ -587,6 +921,7 @@ def _decode_lmdb_worker_chunk( original_keys, frame_format, config, + layout, ) @@ -904,16 +1239,20 @@ def _iter_epoch(self) -> Iterator[list[int]]: return iter(self._sampler) def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: - """Submit one batch as balanced contiguous chunks.""" + """Submit one batch as balanced contiguous chunks. + + The batch layout is resolved here rather than per chunk, so that every + chunk decodes to the same field shapes and the results concatenate. + """ original_keys = self._reader.original_keys(indices) + layout = self._reader.batch_layout(indices) workers = min(self._num_workers, len(original_keys)) base_size, remainder = divmod(len(original_keys), workers) - chunks: list[list[int]] = [] + chunks: list[tuple[list[int], BatchLayout]] = [] start = 0 for worker_index in range(workers): - chunk_size = base_size + int(worker_index < remainder) - stop = start + chunk_size - chunks.append(original_keys[start:stop]) + stop = start + base_size + int(worker_index < remainder) + chunks.append((original_keys[start:stop], layout.chunk(start, stop))) start = stop decode_config = self._reader.worker_decode_config() return [ @@ -923,8 +1262,9 @@ def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: self._reader.frame_format, decode_config, chunk, + chunk_layout, ) - for chunk in chunks + for chunk, chunk_layout in chunks ] def _worth_decoding_in_parallel(self, indices: list[int]) -> bool: @@ -1027,6 +1367,11 @@ def _compute_batch_size(nloc: int, rule: int) -> int: return max(bsi, 1) +#: Seed of the representative shuffle behind :attr:`LmdbDataReader.total_batch`. +#: Fixed so that the reported count is reproducible across calls and processes. +_TOTAL_BATCH_SEED = 0 + + def _parse_positive_rule(spec: str, prefix: str) -> int: """Parse the ``N`` in ``N`` and require ``N > 0``. @@ -1056,16 +1401,17 @@ class LmdbDataReader: Reads LMDB frames and returns dicts of numpy arrays. Backend-specific Dataset classes (PyTorch, JAX, etc.) wrap this. - Datasets are typically mixed-nloc (frames with different atom counts). - The ``mixed_batch`` flag controls batching strategy: + An LMDB typically holds frames of many different atom counts. The + ``batch_size`` rule decides how those frames are grouped: - - ``mixed_batch=False`` (default, old format): each batch contains only - frames with the same nloc. A ``SameNlocBatchSampler`` groups frames - by nloc and yields same-nloc batches. Auto batch_size is computed - per-nloc-group. - - ``mixed_batch=True`` (new format): frames with different nloc can - coexist in one batch (requires padding + mask in collate_fn). - Currently raises ``NotImplementedError`` at collation time. + - Every rule except ``"mix:N"`` keeps a batch homogeneous in atom count, + so no padding is ever needed, whichever layout it is decoded in. + - ``"mix:N"`` allows one batch to span several atom counts. A consumer + reading a flat node axis takes such a batch concatenated; one reading + an ``(nf, nloc, ...)`` axis takes it padded to the batch-wide maximum, + the padded rows carrying ``atype = -1`` so that the neighbor list, the + atomic model and the loss all skip them. The choice is + :meth:`use_ragged_batches`. Parameters ---------- @@ -1090,9 +1436,16 @@ class LmdbDataReader: - ``"filter:N"``: same per-nloc formula as ``"max:N"`` **and** drops every frame whose ``nloc > N`` from the dataset. By construction every retained batch has at most ``N`` atoms. - mixed_batch : bool - If True, allow different nloc in the same batch (future). - If False (default), enforce same-nloc-per-batch. + - ``"mix:N"``: mixed-nloc batching with an atom-axis budget of + ``N``. Frames of different atom counts share a batch, whose atom + axis holds at most ``N`` entries: the total atom count under the + ragged layout, the padded ``nframes * max_nloc`` under the + rectangular one. This is the natural extension of ``"max:N"``: + the bound is on the decoded batch, and a lone frame with + ``nloc > N`` still forms a batch of its own. Filling batches + rather than cutting them at atom-count boundaries also keeps the + frame-level loss terms closer to the weighting an atom budget asks + for than ``"max:N"`` manages; see :func:`_chop_mixed_nloc`. """ def __init__( @@ -1100,12 +1453,10 @@ def __init__( lmdb_path: str, type_map: list[str], batch_size: int | str = "auto", - mixed_batch: bool = False, ) -> None: self.lmdb_path = str(Path(lmdb_path).resolve()) self._type_map = type_map self._env = _open_lmdb(self.lmdb_path) - self.mixed_batch = mixed_batch with self._env.begin() as txn: meta = _read_metadata(txn) @@ -1139,25 +1490,22 @@ def __init__( self._txn = self._env.begin() self._closed = False - # Scan per-frame nloc only when needed for same-nloc batching. - # For mixed_batch=True, skip the scan entirely (future: padding handles it). + # Per-frame atom counts drive every batching rule: same-nloc grouping, + # the ``filter:N`` drop, and the atom-axis layout of a ``mix:N`` batch. # ``orig_frame_nlocs`` / ``orig_frame_system_ids`` are indexed by the # *original* LMDB frame index. After a potential ``filter:N`` drop we # rebuild ``self._frame_nlocs`` / ``self._frame_system_ids`` so they # are parallel arrays over the *dataset* index space (0..len(self)); # the dataset-to-original mapping lives in ``self._retained_keys``. - if not mixed_batch: - # Fast path: use pre-computed frame_nlocs from metadata if available. - # Falls back to scanning each frame's atom_types shape (~10 us/frame). - meta_nlocs = meta.get("frame_nlocs") - if meta_nlocs is not None: - orig_frame_nlocs = [int(n) for n in meta_nlocs] - else: - orig_frame_nlocs = _scan_frame_nlocs( - self._env, self.nframes, self._frame_fmt, self._natoms - ) + # Metadata carries the counts when the writer recorded them; otherwise + # each frame's atom_types shape is scanned (~10 us/frame). + meta_nlocs = meta.get("frame_nlocs") + if meta_nlocs is not None: + orig_frame_nlocs = [int(n) for n in meta_nlocs] else: - orig_frame_nlocs = [] + orig_frame_nlocs = _scan_frame_nlocs( + self._env, self.nframes, self._frame_fmt, self._natoms + ) # Parse frame_system_ids for auto_prob support. ``_nsystems`` must stay # at ``max(original_sid) + 1`` even after filter:N so that user-facing @@ -1171,12 +1519,13 @@ def __init__( orig_frame_system_ids = None self._nsystems = 1 - # Parse batch_size spec. ``auto_rule`` and ``max_rule`` are mutually - # exclusive; ``filter_rule`` implies ``max_rule`` plus dropping frames - # whose nloc exceeds the threshold. + # Parse batch_size spec. ``auto_rule``, ``max_rule`` and ``mix_rule`` + # are mutually exclusive; ``filter_rule`` implies ``max_rule`` plus + # dropping frames whose nloc exceeds the threshold. self._auto_rule: int | None = None self._max_rule: int | None = None self._filter_rule: int | None = None + self._mix_rule: int | None = None if isinstance(batch_size, str): if batch_size == "auto": self._auto_rule = 32 @@ -1187,23 +1536,14 @@ def __init__( elif batch_size.startswith("filter:"): self._filter_rule = _parse_positive_rule(batch_size, "filter:") self._max_rule = self._filter_rule + elif batch_size.startswith("mix:"): + self._mix_rule = _parse_positive_rule(batch_size, "mix:") else: raise ValueError( - f"Unsupported batch_size {batch_size!r}. " - "Expected int, 'auto', 'auto:N', 'max:N', or 'filter:N'." + f"Unsupported batch_size {batch_size!r}. Expected int, " + "'auto', 'auto:N', 'max:N', 'filter:N', or 'mix:N'." ) - # ``filter:N`` needs per-frame nloc to drop oversized frames; the - # ``mixed_batch=True`` fast path skips the nloc scan entirely, so the - # two options are incompatible. Fail fast rather than silently - # retaining every frame and breaking the documented contract. - if self._filter_rule is not None and mixed_batch: - raise ValueError( - "batch_size='filter:N' is incompatible with mixed_batch=True: " - "per-frame nloc is unavailable in the mixed-batch fast path. " - "Use mixed_batch=False, or switch to 'max:N' / a fixed int." - ) - # Determine which original-index frames survive the filter. Without # ``filter:N`` every frame is retained. if self._filter_rule is not None: @@ -1227,12 +1567,9 @@ def __init__( # Re-key _frame_nlocs / _frame_system_ids into the dataset-index # space so that every downstream consumer (nloc_groups, system_groups, - # SameNlocBatchSampler, _expand_indices_by_blocks) operates in a - # single, self-consistent indexing scheme. - if not mixed_batch: - self._frame_nlocs = [orig_frame_nlocs[k] for k in retained_keys] - else: - self._frame_nlocs = [] + # LmdbBatchSampler, _expand_indices_by_blocks) operates in a single, + # self-consistent indexing scheme. + self._frame_nlocs = [orig_frame_nlocs[k] for k in retained_keys] if orig_frame_system_ids is not None: self._frame_system_ids: list[int] | None = [ @@ -1242,12 +1579,12 @@ def __init__( self._frame_system_ids = None # Group retained frames by nloc using dataset indices (0..len-1). - if not mixed_batch: - self._nloc_groups: dict[int, list[int]] = {} - for ds_idx, nloc in enumerate(self._frame_nlocs): - self._nloc_groups.setdefault(nloc, []).append(ds_idx) - else: - self._nloc_groups = {} + # Statistics collection consumes these groups in every batching mode, + # because per-nloc groups are the largest units that stack without + # padding. + self._nloc_groups: dict[int, list[int]] = {} + for ds_idx, nloc in enumerate(self._frame_nlocs): + self._nloc_groups.setdefault(nloc, []).append(ds_idx) # Group retained frames by original system id; the sid numbering is # preserved (no compression) so user-facing auto_prob slices stay @@ -1268,12 +1605,20 @@ def __init__( # valid index domain for __getitem__ is [0, self.nframes). self.nframes = len(retained_keys) - # Default batch_size used only by the index/total_batch estimate. The - # sampler always goes through get_batch_size_for_nloc for real batches. + # Nominal batch size, reported to callers that want a single number. + # The sampler never uses it: same-nloc modes go through + # get_batch_size_for_nloc, and ``mix:N`` sizes each batch by budget. + mean_nloc = ( + sum(self._frame_nlocs) / len(self._frame_nlocs) + if self._frame_nlocs + else self._natoms + ) if self._auto_rule is not None: self.batch_size = _compute_batch_size(self._natoms, self._auto_rule) elif self._max_rule is not None: self.batch_size = max(1, self._max_rule // max(self._natoms, 1)) + elif self._mix_rule is not None: + self.batch_size = max(1, int(self._mix_rule / max(mean_nloc, 1.0))) else: self.batch_size = int(batch_size) @@ -1289,6 +1634,12 @@ def __init__( # Availability signatures are decoded lazily and reused by every # sampler epoch. Registering new requirements invalidates the cache. self._find_signature_cache: dict[int, tuple[tuple[str, bool], ...]] = {} + # Which fields carry an atom axis follows from the requirements, so + # this cache is invalidated alongside the signature cache. + self._per_atom_strides: dict[str, int] | None = None + # Batches are rectangular until a consumer that reads a flat node axis + # asks otherwise; see :meth:`use_ragged_batches`. + self._ragged_batches = False def _resolve_dtype(self, key: str) -> np.dtype: """Resolve the target numpy dtype for a given key. @@ -1335,12 +1686,18 @@ def get_batch_size_for_nloc(self, nloc: int) -> int: - ``filter:N``: same per-nloc formula as ``max:N``; by construction every retained group satisfies ``nloc <= N`` so no overshoot occurs. + - ``mix:N``: ``max(1, floor(N / nloc))``, the count a batch would + hold were every one of its frames this size. Training batches are + sized by budget instead; this value serves the per-nloc statistics + groups, which stack without padding and therefore batch like + ``max:N``. - fixed int: the same value for every nloc group. """ if self._auto_rule is not None: return _compute_batch_size(nloc, self._auto_rule) - if self._max_rule is not None: - return max(1, self._max_rule // max(nloc, 1)) + atom_budget = self._max_rule if self._max_rule is not None else self._mix_rule + if atom_budget is not None: + return max(1, atom_budget // max(nloc, 1)) return self.batch_size def __len__(self) -> int: @@ -1382,14 +1739,127 @@ def original_keys(self, indices: Sequence[int]) -> list[int]: keys.append(self._retained_keys[index]) return keys - def decode_batch(self, indices: Sequence[int]) -> dict[str, Any]: - """Decode a same-nloc batch directly into contiguous NumPy arrays.""" + def batch_pad_nloc(self, indices: Sequence[int]) -> int: + """Return the atom count every frame of one batch is padded to. + + Parameters + ---------- + indices : Sequence[int] + Dataset indices forming one batch. + + Returns + ------- + int + The largest atom count in the batch. Batches drawn from a single + nloc group return that group's atom count, so padding is a no-op. + """ + return max(self._frame_nlocs[int(index)] for index in indices) + + def batch_layout( + self, indices: Sequence[int], *, ragged: bool | None = None + ) -> BatchLayout: + """Return where one batch's per-atom rows belong once decoded. + + The layout fixes the shape of every field of the decoded batch, so a + decode split across worker processes resolves it here once and cuts a + chunk's share from it. Resolving it per chunk would let two chunks + disagree, both on the padded width and on the field classification, + which falls back to comparing a leading axis against the atom count of + whichever frame the chunk happens to start with. + + Parameters + ---------- + indices : Sequence[int] + Dataset indices forming one batch. + ragged : bool, optional + Layout to use, overriding the one configured for training batches. + Consumers with a layout of their own, such as statistics, name it + rather than inherit it. + + Returns + ------- + BatchLayout + The per-frame atom counts, the per-atom strides, and whether the + frames are concatenated or padded to a common width. + """ + return BatchLayout.over( + np.asarray( + [self._frame_nlocs[int(index)] for index in indices], dtype=np.int64 + ), + self.per_atom_strides(), + ragged=self._ragged_batches if ragged is None else ragged, + ) + + @property + def ragged_batches(self) -> bool: + """Whether decoded batches concatenate their frames rather than pad.""" + return self._ragged_batches + + def use_ragged_batches(self, ragged: bool) -> None: + """Select the layout training batches are delivered in. + + The choice belongs to whichever model will consume them: one reading a + flat node axis takes the frames concatenated, one reading an + ``(nf, nloc, ...)`` axis needs them padded to a common width. Only the + trainer sees both the model and the data, so it makes the call, once, + before training starts. Consumers with a layout of their own -- + statistics, validation -- name theirs at the point of use and are + unaffected. + + The layout also decides how the sampler packs frames, since padding is + what makes a batch's cost depend on its widest frame. + + Parameters + ---------- + ragged : bool + Whether to concatenate frames instead of padding them. + """ + self._ragged_batches = ragged + + def per_atom_strides(self) -> dict[str, int]: + """Return the leading-axis entries per atom of each per-atom field. + + Every frame of one LMDB exposes the same fields, so the classification + is a property of the dataset and its registered requirements rather + than of a batch, and is resolved from the first frame once. + + Returns + ------- + dict[str, int] + Entries per atom, keyed by field name. + """ + if self._per_atom_strides is None: + frame = self[0] + self._per_atom_strides = per_atom_strides( + frame, resolve_per_atom_keys(frame, self._decode_config) + ) + return self._per_atom_strides + + def decode_batch( + self, indices: Sequence[int], *, ragged: bool | None = None + ) -> dict[str, Any]: + """Decode one batch directly into contiguous NumPy arrays. + + Parameters + ---------- + indices : Sequence[int] + Dataset indices forming one batch. + ragged : bool, optional + Layout to decode into, overriding the one configured for training + batches. See :meth:`batch_layout`. + + Returns + ------- + dict[str, Any] + One collated batch of contiguous NumPy arrays. + """ self._data_requirements_frozen = True return decode_lmdb_batch( self._transaction(), self.original_keys(indices), self._frame_fmt, self._decode_config, + self.batch_layout(indices, ragged=ragged), ) @property @@ -1397,6 +1867,17 @@ def frame_format(self) -> str: """Format specification used for integer LMDB frame keys.""" return self._frame_fmt + @property + def decode_config(self) -> LmdbDecodeConfig: + """Decoder state for in-process consumers. + + The returned object shares the reader's live requirement mapping and + reading it does not freeze registration, unlike + :meth:`worker_decode_config`, which hands the state to another process + and so must fix it first. + """ + return self._decode_config + def worker_decode_config(self) -> LmdbDecodeConfig: """Freeze and return decoder state for worker serialization.""" self._data_requirements_frozen = True @@ -1418,6 +1899,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N for item in data_requirement: self._data_requirements[item["key"]] = item self._find_signature_cache.clear() + self._per_atom_strides = None def get_find_signature(self, index: int) -> tuple[tuple[str, bool], ...]: """Return the scalar availability signature for one retained frame. @@ -1485,6 +1967,8 @@ def print_summary(self, name: str, prob: Any) -> None: bs_str = f"filter:{self._filter_rule}" elif self._max_rule is not None: bs_str = f"max:{self._max_rule}" + elif self._mix_rule is not None: + bs_str = f"mix:{self._mix_rule}" else: bs_str = str(self.batch_size) @@ -1492,7 +1976,7 @@ def print_summary(self, name: str, prob: Any) -> None: f"LMDB {name}: {self.lmdb_path}, " f"{self.nframes} frames, {n_groups} nloc groups, " f"batch_size={bs_str}, " - f"mixed_batch={self.mixed_batch}" + f"mixed_nloc={self.mixed_nloc}" ) # Print nloc groups in rows of ~10 for readability items = [ @@ -1515,18 +1999,39 @@ def index(self) -> list[int]: @property def total_batch(self) -> int: - if self.mixed_batch: - return math.ceil(self.nframes / self.batch_size) if self.nframes else 0 - total = 0 - for nloc, indices in collect_lmdb_sampling_groups(self): - bs = self.get_batch_size_for_nloc(nloc) - total += (len(indices) + bs - 1) // bs - return total + """Number of batches in one pass over the dataset. + + Every batching rule but ``mix:N`` fixes the count independently of the + order frames are visited in. ``mix:N`` fills batches to an atom budget + instead, so its count follows the shuffle: a pass in dataset order + groups frames of one original system together, which are already close + in atom count, and needs measurably fewer batches than the shuffled + pass training actually performs. The count is therefore taken from a + shuffled pass under a fixed seed, which is drawn from the same + distribution as a training pass while staying reproducible. It remains + an estimate: consult the sampler that will actually be iterated when + the exact count matters, as the trainers do to derive an epoch length. + """ + return len(LmdbBatchSampler(self, shuffle=True, seed=_TOTAL_BATCH_SEED)) @property def batch_sizes(self) -> list[int]: return [self.batch_size] + @property + def mixed_nloc(self) -> bool: + """Whether one batch may span several atom counts.""" + return self._mix_rule is not None + + @property + def atom_budget(self) -> int | None: + """Atom-axis budget of a ``mix:N`` batch, or ``None`` in other modes. + + The axis is measured in the layout the batch will be decoded in: real + atoms when the frames are concatenated, padded slots when they are not. + """ + return self._mix_rule + @property def mixed_type(self) -> bool: """LMDB datasets are always mixed_type (frames may have different compositions).""" @@ -1568,7 +2073,23 @@ def system_nframes(self) -> list[int]: return self._system_nframes -def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: +def _pad_atom_axis(xp: Any, array: Any, length: int, fill: int, device: Any) -> Any: + """Widen one per-atom array's leading axis to ``length``.""" + if array.shape[0] == length: + return array + tail = xp.full( + (length - array.shape[0], *array.shape[1:]), + fill, + dtype=array.dtype, + device=device, + ) + return xp.concat([array, tail], axis=0) + + +def collate_lmdb_frames( + frames: list[dict[str, Any]], + per_atom_keys: frozenset[str] = frozenset(), +) -> dict[str, Any]: """Stack a list of per-frame dicts into a single batch dict. Backend-agnostic via ``array_api_compat``: works for numpy, torch, jax, @@ -1583,6 +2104,23 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: The batch keeps the key order of its frames, which is the order :func:`decode_lmdb_batch` also produces, so a batch is the same mapping whichever of the two decode paths built it. + + Parameters + ---------- + frames : list[dict[str, Any]] + Per-frame dicts to stack, all sharing one label-availability + signature. + per_atom_keys : frozenset[str], optional + Fields whose leading axis is the atom axis. When the frames differ in + atom count these are padded to the batch maximum, with ``atype`` + filled by :data:`PHANTOM_ATOM_TYPE` and the rest zeroed. Leave empty + for uniform batches, where padding would be a no-op anyway. Resolve + the set with :func:`resolve_per_atom_keys`. + + Returns + ------- + dict[str, Any] + One collated batch. """ import array_api_compat @@ -1607,9 +2145,12 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: if any(value != values[0] for value in values[1:]): raise ValueError( f"LMDB batch mixes {key!r} values {values}; " - "SameNlocBatchSampler must group frames by label availability" + "LmdbBatchSampler must group frames by label availability" ) + strides = per_atom_strides(frames[0], per_atom_keys) if per_atom_keys else {} + pad_nloc = max(frame["coord"].shape[0] for frame in frames) if strides else 0 + out: dict[str, Any] = {} for key in frames[0]: if key.startswith("find_"): @@ -1620,6 +2161,12 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: continue elif frames[0][key] is None: out[key] = None + elif key in strides: + length = pad_nloc * strides[key] + fill = _pad_fill_value(key) + out[key] = xp.stack( + [_pad_atom_axis(xp, f[key], length, fill, dev) for f in frames] + ) else: out[key] = xp.stack([f[key] for f in frames]) out["sid"] = xp.asarray([0], dtype=xp.int64, device=dev) @@ -1664,6 +2211,12 @@ def compute_block_targets( stt, end, weight = part.split(":") blocks.append((int(stt), int(end), float(weight))) + # A bare ``prob_sys_size`` names no blocks: it asks for a probability + # proportional to system size, which is what sampling the merged frames + # uniformly already gives. There is nothing to reweight. + if not blocks: + return [] + # Drop blocks that retain zero frames (can happen when ``filter:N`` # eliminates every system in a block). prob_sys_size_ext's per-block # ``nbatch_block / sum(nbatch_block)`` would otherwise propagate NaN @@ -1676,8 +2229,9 @@ def compute_block_targets( ] if not nonempty: log.info( - "compute_block_targets: all blocks are empty in " - f"{auto_prob_style!r}; dataset has no retained frames." + "compute_block_targets: every block of " + f"{auto_prob_style!r} is empty; the dataset retains no frames in " + "any of them, so no reweighting is applied." ) return [] if len(nonempty) < len(blocks): @@ -1886,8 +2440,36 @@ def collect_lmdb_sampling_groups( return groups +def _collect_batch_groups(reader: "LmdbDataReader") -> list[list[int]]: + """Collect the groups a training batch may be drawn from. + + A group is the largest set of frames one batch may span. Label + availability always partitions it, because ``find_*`` flags collapse to a + single scalar per batch. Atom count partitions it as well in every mode + but ``mix:N``, whose decoded batch accommodates unequal counts and so + needs only the availability split. + + Parameters + ---------- + reader : LmdbDataReader + Reader providing the frame grouping and the batching rule. + + Returns + ------- + list[list[int]] + Frame indices per group, in the stable order shared by iteration and + length. + """ + if reader.mixed_nloc: + signature_groups = reader.group_indices_by_find_signature( + list(range(len(reader))) + ) + return [list(signature_groups[key]) for key in sorted(signature_groups)] + return [indices for _nloc, indices in collect_lmdb_sampling_groups(reader)] + + def _allocate_group_block_targets( - groups: list[tuple[int, list[int]]], + groups: list[list[int]], frame_system_ids: list[int] | np.ndarray, block_targets: list[tuple[list[int], int]], ) -> list[list[int]]: @@ -1904,7 +2486,7 @@ def _allocate_group_block_targets( for block_index, (system_ids, _target) in enumerate(block_targets) for system_id in system_ids } - for group_index, (_nloc, indices) in enumerate(groups): + for group_index, indices in enumerate(groups): for index in indices: block_index = system_to_block.get(int(frame_system_ids[index])) if block_index is not None: @@ -1947,24 +2529,120 @@ def _allocate_group_block_targets( return group_targets +def _chop_same_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[int]]: + """Split one homogeneous group into fixed-size batches.""" + batch_size = reader.get_batch_size_for_nloc(reader.frame_nlocs[indices[0]]) + return [ + indices[start : start + batch_size] + for start in range(0, len(indices), batch_size) + ] + + +def _chop_mixed_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[int]]: + """Split one availability group into batches under an atom-axis budget. + + ``mix:N`` budgets the length of a batch's atom axis, and the layout the + batch will be decoded in decides both how that length is measured and the + order the frames are visited in. + + **Ragged.** The frames are concatenated, so the axis is simply their total + atom count and nothing is padded. No frame's cost depends on its + neighbours, so the group is taken in the caller's order, which training has + already shuffled. Sorting would only make each batch homogeneous in system + size, correlating the frames of an optimizer step for no gain. + + **Rectangular.** Every frame is padded to the widest one, so the axis is + ``nframes * max_nloc`` and a batch pays for atoms it does not hold. Sorting + by atom count is what keeps that overhead small: a batch is then a run of + frames adjacent in the sorted order, so its padding is bounded by the + atom-count spread across that run alone. Ties fall back to the caller's + order, so frames of equal atom count still mix freely across epochs. + + Either way the group is cut only where the next frame would push the axis + past the budget, which is the fewest batches obtainable **without + reordering**: a batch's cost does not fall when frames are dropped from its + front, so a batch that starts later can always extend at least as far as + one that starts earlier, and by induction on the batch count, cutting as + late as possible covers the longest prefix for every count. + + Under the rectangular layout the sort is part of the algorithm and the + result is optimal outright, since an exchange argument turns any packing + into contiguous runs of the sorted order. Under the ragged layout the order + is the caller's, so the count is optimal only for that order; reordering + could pack tighter -- the general problem is bin packing -- and is declined + to keep the frames of an optimizer step decorrelated in system size. + Padding, where it exists, is not separately minimized either. + + That the batches are full is what keeps the gradient weighting faithful. A + batch is one optimizer step; its per-atom loss terms pool over the real + labels, so a frame's weight there follows its atom count whatever the + packing, while its frame-level terms (energy, virial) weigh frames equally, + giving a frame the weight ``1 / k_b``. An atom budget makes ``k_b`` follow + the atom count, exactly so under the ragged layout and up to the padding + under the rectangular one, and an under-filled batch raises the weight of + every frame it holds. + + A frame larger than the budget forms a batch of its own, matching how + ``max:N`` treats an oversized nloc group. + + Parameters + ---------- + reader : LmdbDataReader + Provides the per-frame atom counts, the atom budget and the layout. + indices : list[int] + Dataset indices of one label-availability group. + + Returns + ------- + list[list[int]] + Batches whose union is ``indices``. + """ + budget = reader.atom_budget + if budget is None: + raise ValueError("mixed-nloc batching requires a batch_size of 'mix:N'") + + index_array = np.asarray(indices, dtype=np.int64) + nloc_array = np.asarray(reader.frame_nlocs, dtype=np.int64)[index_array] + if not reader.ragged_batches: + order = np.argsort(nloc_array, kind="stable") + index_array, nloc_array = index_array[order], nloc_array[order] + + batches: list[list[int]] = [] + batch_start = 0 + real_atoms = 0 + for position, nloc in enumerate(nloc_array.tolist()): + count = position - batch_start + 1 + # Length of the atom axis this run would occupy once decoded. Under the + # rectangular layout the ascending sort makes the current frame the + # widest, so it alone sets the padded width. + axis = real_atoms + nloc if reader.ragged_batches else count * nloc + if count > 1 and axis > budget: + batches.append(index_array[batch_start:position].tolist()) + batch_start, real_atoms = position, 0 + real_atoms += nloc + batches.append(index_array[batch_start:].tolist()) + return batches + + def _build_all_batches( reader: "LmdbDataReader", shuffle: bool, rng: np.random.Generator, block_targets: list[tuple[list[int], int]] | None = None, ) -> list[list[int]]: - """Build batches homogeneous in atom count and label availability. + """Build batches homogeneous in label availability. - This is the shared batch-construction logic used by both - SameNlocBatchSampler (single-GPU) and DistributedSameNlocBatchSampler. + Groups are chopped into batches, then interleaved round-robin so that + consecutive batches come from different groups. Under ``mix:N`` a batch + also spans several atom counts; every other mode keeps it uniform. Parameters ---------- reader : LmdbDataReader - Provides nloc_groups and get_batch_size_for_nloc. + Provides the frame grouping and the batching rule. shuffle : bool - Whether to shuffle indices within each nloc group and - shuffle the final batch order. + Whether to shuffle indices within each group and shuffle the final + batch order. rng : np.random.Generator Random number generator (deterministic for reproducibility). block_targets : list[tuple[list[int], int]] or None @@ -1974,9 +2652,10 @@ def _build_all_batches( Returns ------- list[list[int]] - Each inner list has one nloc and one scalar ``find_*`` signature. + Each inner list has one scalar ``find_*`` signature. """ - groups = collect_lmdb_sampling_groups(reader) + groups = _collect_batch_groups(reader) + chop = _chop_mixed_nloc if reader.mixed_nloc else _chop_same_nloc # Build per-group batches group_batches: list[list[list[int]]] = [] @@ -2001,7 +2680,7 @@ def _build_all_batches( for sid, blk in sys_to_block.items(): sid_to_blk_arr[sid] = blk - for group_index, (nloc, original_indices) in enumerate(groups): + for group_index, original_indices in enumerate(groups): indices = original_indices # Expand each availability group independently using targets that # were allocated globally, preserving both scalar flags and totals. @@ -2016,11 +2695,7 @@ def _build_all_batches( ) if shuffle: rng.shuffle(indices) - bs = reader.get_batch_size_for_nloc(nloc) - batches = [] - for start in range(0, len(indices), bs): - batches.append(indices[start : start + bs]) - group_batches.append(batches) + group_batches.append(chop(reader, indices) if indices else []) # Interleave groups round-robin all_batches: list[list[int]] = [] @@ -2037,26 +2712,31 @@ def _build_all_batches( return all_batches -class SameNlocBatchSampler: - """Batch sampler that groups frames by nloc and ``find_*`` signature. +class LmdbBatchSampler: + """Batch sampler over an LMDB, grouped by ``find_*`` signature. - For mixed-nloc datasets with mixed_batch=False: each batch contains only - frames with the same nloc and label availability. Within each group, - frames are shuffled. Groups are interleaved round-robin so training sees - diverse nloc and label combinations. + Every batch carries one label-availability signature, because ``find_*`` + flags collapse to a single scalar per batch. Atom count is handled by the + reader's batching rule: all rules but ``mix:N`` additionally keep a batch + uniform in atom count, while ``mix:N`` fills batches to the atom budget + its decoded layout is measured against. Groups are interleaved round-robin + and the batch order is then shuffled, so training sees a varied mix. - When auto batch_size is used, batch_size is computed per-nloc-group. - - The sampler is deterministic for a fixed seed and epoch. Use - :meth:`set_epoch` to select a different reproducible sequence for each - training pass. + The sampler serves one pass at a time, drawn from ``seed + epoch``. The + pending pass is materialized before it is served, which is what lets + ``__len__`` report exactly what ``__iter__`` will yield: under ``mix:N`` + the batch count follows the shuffle, because batches are filled to an atom + budget rather than to a fixed frame count. Serving a pass advances the + epoch, so a caller that just re-iterates sees a different shuffle every + time, and :meth:`set_epoch` repositions that progression for a caller -- + a distributed run, a resumed one -- that needs to name the pass instead. Parameters ---------- reader : LmdbDataReader - The dataset reader (provides nloc_groups, get_batch_size_for_nloc). + The dataset reader providing the frame grouping and batching rule. shuffle : bool - Whether to shuffle within each nloc group each epoch. + Whether to shuffle within each group and shuffle the batch order. seed : int or None Random seed for reproducibility. block_targets : list[tuple[list[int], int]] or None @@ -2075,61 +2755,71 @@ def __init__( self._seed = seed self._epoch = 0 self._block_targets = block_targets + self._batches: list[list[int]] | None = None + + def batches(self) -> list[list[int]]: + """Return the batch list of the pending pass, building it if needed. + + Returns + ------- + list[list[int]] + Dataset indices grouped into batches, one scalar ``find_*`` + signature each. + """ + if self._batches is None: + seed = None if self._seed is None else self._seed + self._epoch + self._batches = _build_all_batches( + self._reader, + self._shuffle, + np.random.default_rng(seed), + self._block_targets, + ) + return self._batches def set_epoch(self, epoch: int) -> None: - """Set the epoch used to derive the deterministic shuffle state. + """Select the pass to serve, discarding any pass still pending. Parameters ---------- epoch : int Zero-based training epoch. """ - self._epoch = epoch + if epoch != self._epoch: + self._epoch = epoch + self._batches = None + + def refresh_batch_count(self) -> None: + """Discard the pending pass after the frame grouping changed. + + The pass is materialized ahead of iteration so that ``__len__`` can + report it exactly, which leaves it stale once new data requirements + repartition the frames by label availability. + """ + self._batches = None def __iter__(self) -> Iterator[list[int]]: - """Yield batches of frame indices, all with the same nloc.""" - seed = None if self._seed is None else self._seed + self._epoch - rng = np.random.default_rng(seed) - yield from _build_all_batches( - self._reader, self._shuffle, rng, self._block_targets - ) + """Yield the pending pass, and move the epoch on to its successor.""" + batches = self.batches() + self.set_epoch(self._epoch + 1) + yield from batches def __len__(self) -> int: - """Total batches across nloc and label-availability groups.""" - groups = collect_lmdb_sampling_groups(self._reader) - group_block_targets = None - assigned_system_ids: set[int] = set() - if self._block_targets and self._reader.frame_system_ids is not None: - group_block_targets = _allocate_group_block_targets( - groups, - self._reader.frame_system_ids, - self._block_targets, - ) - assigned_system_ids = { - system_id - for system_ids, _target in self._block_targets - for system_id in system_ids - } - - total = 0 - for group_index, (nloc, indices) in enumerate(groups): - bs = self._reader.get_batch_size_for_nloc(nloc) - n = len(indices) - if ( - group_block_targets is not None - and self._reader.frame_system_ids is not None - ): - unassigned = sum( - int(self._reader.frame_system_ids[index]) not in assigned_system_ids - for index in indices - ) - n = unassigned + sum(group_block_targets[group_index]) - total += (n + bs - 1) // bs - return total + """Number of batches the pending pass holds.""" + return len(self.batches()) + + @property + def total_batches(self) -> int: + """Number of batches the pending pass holds over the whole dataset. + + The same count as ``len(self)`` here, and the two part company only in + the distributed sampler, so a caller after a global figure need not + know which of the two it holds. + """ + return len(self) -class DistributedSameNlocBatchSampler: - """Distributed wrapper for same-nloc batch sampling. +class DistributedLmdbBatchSampler: + """Distributed wrapper for LMDB batch sampling. All ranks build the same deterministic global batch list (using ``seed + epoch``). The list is padded deterministically when its length is @@ -2144,8 +2834,7 @@ class DistributedSameNlocBatchSampler: Parameters ---------- reader : LmdbDataReader - The dataset reader (provides nloc_groups, get_batch_size_for_nloc, - frame_nlocs). + The dataset reader providing the frame grouping and batching rule. rank : int Rank of the current process. world_size : int @@ -2174,17 +2863,11 @@ def __init__( self._seed = seed if seed is not None else 0 self._epoch = 0 self._block_targets = block_targets - self.refresh_batch_count() + self._global: LmdbBatchSampler | None = None def refresh_batch_count(self) -> None: - """Refresh the cached global count after sampling groups change.""" - self._total_batches = len( - SameNlocBatchSampler( - self._reader, - shuffle=False, - block_targets=self._block_targets, - ) - ) + """Discard the pending global pass after the frame grouping changed.""" + self._global = None def set_epoch(self, epoch: int) -> None: """Set epoch for deterministic cross-rank shuffling. @@ -2192,31 +2875,38 @@ def set_epoch(self, epoch: int) -> None: Call this before each training epoch/cycle to get different but reproducible batch orderings across epochs. """ - self._epoch = epoch + if epoch != self._epoch: + self._epoch = epoch + self._global = None + + def _global_batches(self) -> list[list[int]]: + """Return the batch list every rank builds identically.""" + if self._global is None: + self._global = LmdbBatchSampler( + self._reader, + shuffle=self._shuffle, + seed=self._seed + self._epoch, + block_targets=self._block_targets, + ) + return self._global.batches() def __iter__(self) -> Iterator[list[int]]: """Yield this rank's partition of the global batch list.""" - # All ranks build the same global batch list deterministically - rng = np.random.default_rng(self._seed + self._epoch) - all_batches = _build_all_batches( - self._reader, self._shuffle, rng, self._block_targets - ) - # Partition to this rank - yield from self._partition_batches(all_batches) + yield from self._partition_batches(self._global_batches()) def _partition_batches(self, all_batches: list[list[int]]) -> list[list[int]]: """Partition global batches to this rank. The default pads the global list to a multiple of ``world_size`` and then takes ``all_batches[rank::world_size]``. This gives good nloc - diversity per rank since batches are interleaved across nloc groups - before shuffling, while ensuring that every rank yields the same - number of batches. + diversity per rank since batches are interleaved across groups before + shuffling, while ensuring that every rank yields the same number of + batches. Override this method for custom load-balancing. For example, a greedy algorithm could assign batches to ranks based on estimated - compute cost (``reader.frame_nlocs[batch[0]]`` gives the nloc of - each batch). + compute cost (``reader.batch_pad_nloc(batch) * len(batch)`` gives the + padded cost of each batch). """ if not all_batches: return [] @@ -2233,12 +2923,12 @@ def _partition_batches(self, all_batches: list[list[int]]) -> list[list[int]]: def __len__(self) -> int: """Number of batches for this rank.""" - return (self._total_batches + self._world_size - 1) // self._world_size + return len(self._partition_batches(self._global_batches())) @property def total_batches(self) -> int: - """Return the global batch count before distributed padding.""" - return self._total_batches + """Number of batches one full pass holds, before the per-rank split.""" + return len(self._global_batches()) @property def rank(self) -> int: diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index c047d37e97..32c0905a2d 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -44,6 +44,8 @@ GraphLayout, NeighborGraph, apply_pair_exclusion, + compact_nodes, + expand_node_values, frame_id_from_n_node, node_ownership_mask, node_validity_mask, @@ -75,8 +77,10 @@ "build_neighbor_graph_ase", "canonicalize_neighbor_graph", "center_edge_pairs", + "compact_nodes", "edge_env_mat", "edge_force_virial", + "expand_node_values", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py index 8c7288c7b4..42fae432a8 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -169,7 +169,7 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: coord, box, nframe_all, - nloc, + np.full(nf, nloc, dtype=np.int64), layout=layout, ) if pair_excl is not None: diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py index d8c8f73c27..b210ef48de 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py @@ -3,7 +3,7 @@ ``neighbor_graph_from_ijs`` is the canonical sparse converter: it takes an already-built sparse edge list -- per-edge center ``i``, neighbor ``j`` (both -per-frame LOCAL indices in ``[0, nloc)``) and integer periodic-image shift ``S`` +indices within their own frame) and integer periodic-image shift ``S`` -- and emits a :class:`NeighborGraph` whose ``edge_vec`` is recomputed DIFFERENTIABLY from ``coord``/``box`` (it never trusts the builder's distance vectors). It is the format-conversion step shared by every O(N) search backend @@ -46,7 +46,7 @@ def neighbor_graph_from_ijs( coord: Array, box: Array | None, nframe_id: Array, - nloc: int, + n_node: Array, layout: GraphLayout | None = None, *, with_csr: bool = False, @@ -61,19 +61,22 @@ def neighbor_graph_from_ijs( Parameters ---------- i - (E,) int per-edge center, per-frame LOCAL index in ``[0, nloc)``. + (E,) int per-edge center, index within its own frame. j - (E,) int per-edge neighbor, per-frame LOCAL index in ``[0, nloc)``. + (E,) int per-edge neighbor, index within its own frame. S (E, 3) int periodic-image shift: the neighbor sits at ``coord[j] + S @ box``. coord - (nf, nloc, 3) local coordinates. + (N, 3) local coordinates, frame-major over ``n_node``. box (nf, 3, 3) simulation cell, or ``None`` for non-periodic (``S`` ignored). nframe_id (E,) int frame index of each edge. - nloc - number of local atoms per frame (used for the frame-major node offset). + n_node + (nf,) int atoms per frame. Frames occupy contiguous blocks of the node + axis in order, so the prefix sums of this vector are the frame offsets + that turn a within-frame index into a node index. A batch padded to a + common width is the special case where every entry is that width. layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. with_csr @@ -86,24 +89,24 @@ def neighbor_graph_from_ijs( Returns ------- NeighborGraph - ``edge_index = [j + nframe_id*nloc, i + nframe_id*nloc]`` (src=neighbor, - dst=center); ``edge_vec = coord[j] + S@box - coord[i]``; ``n_node`` is - ``nloc`` per frame. + ``edge_index`` holds node indices (src=neighbor, dst=center) and + ``edge_vec = coord[j] + S@box - coord[i]``. """ if layout is None: layout = GraphLayout() with_csr = with_csr or canonicalize xp = array_api_compat.array_namespace(coord) dev = array_api_compat.device(coord) - nf = coord.shape[0] - coord = xp.reshape(coord, (nf, nloc, 3)) + n_node = xp.astype(xp.asarray(n_node, device=dev), xp.int64) + nf = n_node.shape[0] + coord_flat = xp.reshape(coord, (-1, 3)) i = xp.astype(xp.asarray(i, device=dev), xp.int64) j = xp.astype(xp.asarray(j, device=dev), xp.int64) nframe_id = xp.astype(xp.asarray(nframe_id, device=dev), xp.int64) - # flat frame-major node indices - i_flat = i + nframe_id * nloc - j_flat = j + nframe_id * nloc - coord_flat = xp.reshape(coord, (nf * nloc, 3)) + # Within-frame indices become node indices through the frame offsets. + offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0) + i_flat = i + offset + j_flat = j + offset r_i = xp.take(coord_flat, i_flat, axis=0) r_j = xp.take(coord_flat, j_flat, axis=0) edge_vec = r_j - r_i @@ -120,7 +123,6 @@ def neighbor_graph_from_ijs( edge_index, edge_vec, edge_mask = pad_and_guard_edges( edge_index, edge_vec, layout.edge_capacity, layout.min_edges ) - n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) if not with_csr: return NeighborGraph( n_node=n_node, @@ -140,7 +142,7 @@ def neighbor_graph_from_ijs( edge_index, edge_vec, edge_mask, - nf * nloc, + int(coord_flat.shape[0]), canonicalize=canonicalize, ) return NeighborGraph( diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index f702c2c6a1..baeef186ef 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -396,6 +396,129 @@ def apply_pair_exclusion( return out +def compact_nodes( + graph: NeighborGraph, node_mask: Array +) -> tuple[NeighborGraph, Array]: + """Renumber a graph onto the subset of its nodes selected by a mask. + + A batch whose frames hold unequal atom counts reaches the model as a + rectangular tensor padded with phantom atoms, and the graph builders lay + their node axis out over that padded shape. Every builder already refuses + an edge touching a phantom, so the phantoms survive as isolated nodes + whose only effect is to make the network evaluate them. Dropping them here + leaves each frame occupying exactly as many nodes as it has real atoms, + which is the layout ``n_node`` was always able to express. + + Node order is preserved, so each frame keeps one contiguous block and the + frame-major invariant every per-frame reduction relies on still holds. + + Parameters + ---------- + graph : NeighborGraph + Graph whose node axis is to be compacted. It must carry no edge + incident on a masked-out node, and no multi-rank halo split. + node_mask : Array + Boolean mask over the flat node axis with shape ``(N,)``, ``True`` for + the nodes to retain. + + Returns + ------- + NeighborGraph + The graph over the retained nodes, with ``n_node`` recounted and + ``edge_index`` renumbered. CSR views, when present, are rebuilt. + Array + Positions of the retained nodes on the original axis, shape + ``(n_kept,)``. Gathering a per-node tensor with it moves that tensor + onto the compacted axis; :func:`expand_node_values` inverts it. + + Raises + ------ + ValueError + If the graph carries a halo split, or if some edge is incident on a + node the mask drops. + """ + import dataclasses + + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + segment_sum, + ) + + if graph.n_local is not None: + raise ValueError("cannot compact the node axis of a local-plus-halo graph") + xp = array_api_compat.array_namespace(graph.n_node, node_mask) + device = array_api_compat.device(graph.n_node) + n_total = node_mask.shape[0] + + keep_index = xp.reshape(xp.nonzero(node_mask)[0], (-1,)) + # Position of each retained node on the compacted axis, and -1 for the + # nodes that go away. Renumbering by a prefix sum keeps the frame blocks + # contiguous and in order. + renumber = xp.cumulative_sum(xp.astype(node_mask, xp.int64)) - 1 + renumber = xp.where(node_mask, renumber, xp.asarray(-1, device=device)) + + frame_id = frame_id_from_n_node(graph.n_node, n_total=n_total) + n_node = xp.astype( + segment_sum(xp.astype(node_mask, xp.int64), frame_id, graph.n_node.shape[0]), + graph.n_node.dtype, + ) + + edge_index = xp.take(renumber, xp.reshape(graph.edge_index, (-1,)), axis=0) + edge_index = xp.reshape(edge_index, graph.edge_index.shape) + # A dropped node with an edge would leave -1 behind on a real edge, which + # would silently address the last node of the compacted axis downstream. + if bool(xp.any(xp.logical_and(edge_index < 0, graph.edge_mask[None, :]))): + raise ValueError("cannot compact a node that still carries an edge") + edge_index = xp.astype( + xp.maximum(edge_index, xp.asarray(0, device=device)), graph.edge_index.dtype + ) + + compacted = dataclasses.replace( + graph, + n_node=n_node, + edge_index=edge_index, + destination_order=None, + destination_row_ptr=None, + source_order=None, + source_row_ptr=None, + destination_sorted=False, + ) + if graph.destination_row_ptr is not None: + from deepmd.dpmodel.utils.neighbor_graph.csr import ( + attach_edge_csr, + ) + + compacted = attach_edge_csr(compacted, int(keep_index.shape[0])) + return compacted, keep_index + + +def expand_node_values(values: Array, keep_index: Array, n_total: int) -> Array: + """Scatter a compacted per-node tensor back onto a padded node axis. + + Inverse of the gather that :func:`compact_nodes` describes. Positions the + compaction dropped read as zero, which is what the dropped nodes + contributed to every physical quantity in the first place. + + Parameters + ---------- + values : Array + Per-node tensor on the compacted axis with shape ``(n_kept, ...)``. + keep_index : Array + Positions of the retained nodes, as returned by :func:`compact_nodes`. + n_total : int + Size of the padded node axis to scatter onto. + + Returns + ------- + Array + Tensor of shape ``(n_total, ...)``. + """ + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + segment_sum, + ) + + return segment_sum(values, keep_index, n_total) + + def node_validity_mask(n_node: Array, n_total: int) -> Array: """Derive the (n_total,) real-vs-padding node mask from per-frame counts. diff --git a/deepmd/pt/loss/dens.py b/deepmd/pt/loss/dens.py index 03e1c297e4..5c6fb3edb5 100644 --- a/deepmd/pt/loss/dens.py +++ b/deepmd/pt/loss/dens.py @@ -4,7 +4,6 @@ ) import torch -import torch.nn.functional as F from deepmd.pt.loss.ener import ( EnergyStdLoss, @@ -164,9 +163,18 @@ def _prepare_dens_inputs( torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, bool, ]: - """Build noisy coordinates and mixed targets for one forward pass.""" + """Build noisy coordinates and mixed targets for one forward pass. + + Returns the corrupted and clean atom subsets as two masks. Both are + restricted to real atoms: a mixed-nloc batch is padded to a rectangular + shape with phantom slots that carry no physical site, and a phantom is + neither corrupted nor clean. The two masks are built independently + rather than as complements of one another, since complements over the + padded width would hand every phantom to the clean subset. + """ atype = input_dict["atype"] nf, nloc = atype.shape[:2] coord_raw = input_dict["coord"] @@ -185,22 +193,22 @@ def _prepare_dens_inputs( ).item() < self.dens_prob ) + real_mask = atype >= 0 noise_mask = torch.zeros((nf, nloc), dtype=torch.bool, device=coord.device) noise_vec = torch.zeros_like(coord) if use_dens: if self.dens_corrupt_ratio is None: - noise_mask = torch.ones( - (nf, nloc), dtype=torch.bool, device=coord.device - ) + noise_mask = real_mask else: noise_mask = ( torch.rand( (nf, nloc), dtype=GLOBAL_PT_FLOAT_PRECISION, device=coord.device ) < self.dens_corrupt_ratio - ) + ) & real_mask noise_vec = torch.randn_like(coord) * self.dens_std noise_vec = noise_vec * noise_mask.unsqueeze(-1) + clean_mask = real_mask & ~noise_mask coord_model = coord + noise_vec # DeNS predicts normalized noise epsilon / sigma for corrupted atoms. @@ -214,7 +222,14 @@ def _prepare_dens_inputs( model_input["noise_mask"] = noise_mask if use_dens: model_input["force_input"] = force_label - return model_input, force_label, noise_target, noise_mask, use_dens + return ( + model_input, + force_label, + noise_target, + noise_mask, + clean_mask, + use_dens, + ) @staticmethod def _get_sezm_atomic_model(model: torch.nn.Module) -> Any: @@ -270,8 +285,14 @@ def forward( learning_rate: float, mae: bool = False, ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: - """Return loss on SeZM `dens` energy and direct-force/noise outputs.""" - model_input, force_label, noise_target, noise_mask, use_dens = ( + """Return loss on SeZM `dens` energy and direct-force/noise outputs. + + ``natoms`` is the padded width of the batch and is superseded here by + the per-frame real atom count read off ``atype``; it remains in the + signature to satisfy the shared loss interface. + """ + del natoms + model_input, force_label, noise_target, noise_mask, clean_mask, use_dens = ( self._prepare_dens_inputs( input_dict, label, @@ -288,7 +309,14 @@ def forward( loss = force_label.new_zeros((), dtype=env.GLOBAL_PT_FLOAT_PRECISION) more_loss: dict[str, torch.Tensor] = {} - atom_norm = 1.0 / natoms + # The energy terms are extensive, so a frame's residual is divided by + # that frame's own real atom count; a mixed-nloc batch pads to a common + # width with phantom atoms that belong to no system. The factor is + # therefore applied inside the mean over frames, where on a batch of + # uniform atom count it reduces to the scalar ``1 / natoms``. + inv_natoms = 1.0 / (input_dict["atype"] >= 0).sum(dim=-1, keepdim=True).to( + dtype=force_label.dtype + ) # [nf, 1] if self.has_e and "energy" in model_pred and "energy" in label: energy_pred = model_pred.get("energy_norm", model_pred["energy"]) @@ -318,24 +346,23 @@ def forward( l2_ener_loss.detach(), find_energy, ) - loss += atom_norm * (pref_e * l2_ener_loss) - rmse_e = ( - torch.mean(torch.square(energy_pred_phys - energy_label)).sqrt() - * atom_norm + loss += pref_e * torch.mean( + torch.square(energy_pred - energy_label_norm) * inv_natoms ) + rmse_e = torch.mean( + torch.square(energy_pred_phys - energy_label) * inv_natoms**2 + ).sqrt() more_loss["rmse_e"] = self.display_if_exist( rmse_e.detach(), find_energy, ) elif self.loss_func == "mae": - l1_ener_loss = F.l1_loss( - energy_pred.reshape(-1), - energy_label_norm.reshape(-1), - reduction="mean", + l1_ener_loss = torch.mean( + torch.abs(energy_pred - energy_label_norm) * inv_natoms ) - loss += atom_norm * (pref_e * l1_ener_loss) - mae_e = ( - torch.mean(torch.abs(energy_pred_phys - energy_label)) * atom_norm + loss += pref_e * l1_ener_loss + mae_e = torch.mean( + torch.abs(energy_pred_phys - energy_label) * inv_natoms ) more_loss["mae_e"] = self.display_if_exist( mae_e.detach(), @@ -346,8 +373,8 @@ def forward( f"Loss type {self.loss_func} is not implemented for `dens` energy loss." ) if mae: - mae_e = ( - torch.mean(torch.abs(energy_pred_phys - energy_label)) * atom_norm + mae_e = torch.mean( + torch.abs(energy_pred_phys - energy_label) * inv_natoms ) more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy) mae_e_all = torch.mean(torch.abs(energy_pred_phys - energy_label)) @@ -386,10 +413,17 @@ def forward( else: force_pred_phys = atomic_model.denorm_dens_force(clean_force_pred_norm) force_target_norm = atomic_model.norm_dens_force(force_label) - clean_mask = ~noise_mask noise_only_mask = noise_mask if use_dens else torch.zeros_like(noise_mask) - clean_fraction = clean_mask.to(dtype=GLOBAL_PT_FLOAT_PRECISION).mean() - noise_fraction = noise_only_mask.to(dtype=GLOBAL_PT_FLOAT_PRECISION).mean() + # Each subset loss averages over its own atoms and is then scaled + # by that subset's share, so the two recombine into one average + # over the real atoms. The share is taken against the real-atom + # count rather than the padded width, or phantom slots would + # dilute both terms by the padding fraction. + real_count = ( + (input_dict["atype"] >= 0).sum().to(dtype=GLOBAL_PT_FLOAT_PRECISION) + ) + clean_fraction = clean_mask.sum() / real_count + noise_fraction = noise_only_mask.sum() / real_count clean_force_loss = self._compute_force_subset_loss( clean_force_pred_norm[clean_mask].reshape(-1, 3), force_target_norm[clean_mask].reshape(-1, 3), diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 7046f73687..8f24e94d22 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -157,7 +157,7 @@ def forward( local_tensor_pred_dos - local_tensor_label_dos ) # [nf, natoms, numb_dos] if "mask" in model_pred: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=numb_dos). maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] l2_local_loss_dos = masked_atom_mean( torch.square(diff), maskf, self.numb_dos @@ -186,7 +186,7 @@ def forward( local_tensor_pred_cdf - local_tensor_label_cdf ) # [nf, natoms, numb_dos] if "mask" in model_pred: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=numb_dos). maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] l2_local_loss_cdf = masked_atom_mean( torch.square(diff), maskf, self.numb_dos diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 468b95f1eb..381d07d97c 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -402,7 +402,6 @@ def forward( if maskf is not None: # Idiom 1 (per-atom masked mean, ncomp=3). diff_f_3d = diff_f.reshape(_nf, _nloc, 3) - maskf_col = maskf.reshape(_nf, _nloc, 1) # Masked MSE computed for rmse_f display regardless of use_huber. l2_f_masked = masked_atom_mean( torch.square(diff_f_3d), maskf, 3 @@ -410,6 +409,12 @@ def forward( if not self.use_huber: loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: + # ``f_use_norm`` selects the residual an atom + # contributes: three independent components, or the + # single L2 norm of its force-error vector. That + # choice sets the label count per atom, which is + # exactly the ``ncomp`` the pooled reduction + # divides by. if not self.f_use_norm: abs_e = torch.abs(diff_f_3d) quad = 0.5 * torch.square(diff_f_3d) @@ -418,9 +423,8 @@ def forward( ) huber_elem = torch.where( abs_e <= self._huber_delta_force, quad, lin - ) - huber_masked = huber_elem * maskf_col - per_frame_dof = maskf.sum(dim=-1) * 3 + ) # [nf, nloc, 3] + huber_ncomp = 3 else: diff_3 = (force_label - force_pred).reshape( _nf, _nloc, 3 @@ -433,13 +437,13 @@ def forward( lin_n = self._huber_delta_force * ( abs_n - 0.5 * self._huber_delta_force ) - huber_n = torch.where( + huber_elem = torch.where( abs_n <= self._huber_delta_force, quad_n, lin_n - ) - huber_masked = (huber_n * maskf).reshape(_nf, _nloc, 1) - per_frame_dof = maskf.sum(dim=-1) - per_frame_sum = huber_masked.reshape(_nf, -1).sum(dim=-1) - l_huber_masked = torch.mean(per_frame_sum / per_frame_dof) + ).reshape(_nf, _nloc, 1) + huber_ncomp = 1 + l_huber_masked = masked_atom_mean( + huber_elem, maskf, huber_ncomp + ) loss += pref_f * l_huber_masked else: if not self.use_huber: @@ -486,10 +490,10 @@ def forward( norm_2d = torch.linalg.vector_norm( diff_3.reshape(-1, 3), ord=2, dim=1 ).reshape(_nf, _nloc) - masked_norm = norm_2d * maskf - per_frame_sum = masked_norm.sum(dim=-1) - per_frame_dof = maskf.sum(dim=-1) - l1_f_masked = torch.mean(per_frame_sum / per_frame_dof) + # One L2 norm per atom, hence one label per atom. + l1_f_masked = masked_atom_mean( + norm_2d.reshape(_nf, _nloc, 1), maskf, 1 + ) more_loss["mae_f"] = self.display_if_exist( l1_f_masked.detach(), find_force ) @@ -726,7 +730,6 @@ def forward( # Idiom 1 (per-atom masked mean, ncomp=1). ae_2d = atom_ener.reshape(_nf, _nloc) ae_hat_2d = atom_ener_label.reshape(_nf, _nloc) - per_frame_dof = maskf.sum(dim=-1) # [nf], kept for huber branch l2_ae_masked = masked_atom_mean( torch.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 ) @@ -742,8 +745,7 @@ def forward( huber_ae = torch.where( abs_ae <= self._huber_delta_energy, quad_ae, lin_ae ) - huber_ae_m = huber_ae * maskf - l_huber_ae = torch.mean(huber_ae_m.sum(dim=-1) / per_frame_dof) + l_huber_ae = masked_atom_mean(huber_ae[:, :, None], maskf, 1) loss += pref_ae * l_huber_ae rmse_ae = l2_ae_masked.sqrt() more_loss["rmse_ae"] = self.display_if_exist( diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 660cbd0337..fa324315e7 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -133,7 +133,7 @@ def forward( ) diff = diff * atomic_weight if "mask" in model_pred: - # idiom 1: per-frame masked mean, then average over frames + # Idiom 1 (per-atom masked mean, ncomp=tensor_size). maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] diff3d = diff.reshape( local_tensor_pred.shape[0], natoms, self.tensor_size diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index f7c8337433..0d907f1c6c 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -770,6 +770,20 @@ def __init__( # Forward Methods # ========================================================================= + def _sanitize_atom_types( + self, + atype: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return lookup-safe atom types and the physical-atom mask. + + Phantom atoms use a negative type sentinel for rectangular mixed-nloc + padding. Every type-table lookup receives the safe tensor, while the + mask keeps the substituted row from acquiring physical meaning. + """ + real_atom = self.atomic_model.make_atom_mask(atype) + safe_atype = torch.where(real_atom, atype, torch.zeros_like(atype)) + return safe_atype, real_atom + def forward( self, coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], @@ -1497,9 +1511,30 @@ def core_compute( ): spin = spin.detach().requires_grad_(True) - descriptor_atype = extended_atype if comm_dict is not None else atype - if descriptor_atype is None: + # === Atom mask === + # Phantom atoms (atype < 0) pad a mixed-nloc batch to a rectangular + # shape and stand for no physical site. The edge builders keep them out + # of every edge, so the networks see them as isolated nodes, and the + # mask below zeroes whatever those nodes produce. Their type is still + # read as a table index by the type embedding, the fitting bias and the + # exclusion masks, none of which has a row for it. ``atype`` is + # therefore rebound to a sanitized copy for the remainder of this + # method, which is where a phantom would otherwise reach such a table. + # The substitute type is immaterial: it reaches no real atom through + # any edge, and the mask discards everything it produces. + atype, real_atom = self._sanitize_atom_types(atype) + atom_mask = real_atom.to(torch.int32) + if self.atomic_model.atom_excl is not None: + atom_mask = atom_mask * self.atomic_model.atom_excl(atype) + + if comm_dict is None: + descriptor_atype = atype + elif extended_atype is None: raise ValueError("`extended_atype` is required with `comm_dict`.") + else: + # Ghost atoms carry the type of the local atom they image, so this + # sanitizes the phantoms among them on the same grounds as above. + descriptor_atype, _ = self._sanitize_atom_types(extended_atype) inter_potential_edge_mask = self._make_inter_potential_edge_mask( descriptor_atype, edge_index, @@ -1526,11 +1561,6 @@ def core_compute( nloc=nloc, ) - # === Atom mask === - atom_mask = self.atomic_model.make_atom_mask(atype).to(torch.int32) - if self.atomic_model.atom_excl is not None: - atom_mask = atom_mask * self.atomic_model.atom_excl(atype) - # === Step 3. Fitting net === # The same fitting forward serves both modes; ``embedding_only`` only asks # it to also return the last hidden activation. @@ -1555,9 +1585,14 @@ def core_compute( structural_feature = ( atomic_feature * atom_mask[:, :, None].to(atomic_feature.dtype) ).sum(dim=1) + # The per-atom outputs are zeroed on the phantom rows alone, so + # that a padded batch reports nothing for a slot holding no atom. + # Excluded atoms, which ``atom_mask`` also covers, keep their + # per-atom embeddings and are dropped from the pooled sum only. + real = real_atom[:, :, None] return { - "descriptor": descriptor, - "atomic_feature": atomic_feature, + "descriptor": descriptor * real.to(descriptor.dtype), + "atomic_feature": atomic_feature * real.to(atomic_feature.dtype), "structural_feature": structural_feature, } @@ -3239,8 +3274,7 @@ def _make_inter_potential_edge_mask( atom_excl = self.atomic_model.atom_excl if atom_excl is not None: - atom_is_present = self.atomic_model.make_atom_mask(atype) - safe_atype = torch.where(atom_is_present, atype, 0) + safe_atype, atom_is_present = self._sanitize_atom_types(atype) atom_is_included = atom_is_present & atom_excl(safe_atype).to(torch.bool) atom_is_included = atom_is_included.reshape(-1) keep = ( diff --git a/deepmd/pt/model/model/sezm_native_spin_model.py b/deepmd/pt/model/model/sezm_native_spin_model.py index aeec0ae74e..9cc14355b9 100644 --- a/deepmd/pt/model/model/sezm_native_spin_model.py +++ b/deepmd/pt/model/model/sezm_native_spin_model.py @@ -84,6 +84,15 @@ def __init__( # Forward Methods # ========================================================================= + def _make_mask_mag(self, atype: torch.Tensor) -> torch.Tensor: + """Return the magnetic-type mask without indexing phantom sentinels.""" + safe_atype, real_atom = self._sanitize_atom_types(atype) + spin_active = self.spin_mask.index_select( + 0, + safe_atype.reshape(-1), + ).reshape(*atype.shape, 1) + return (spin_active > 0.0) & real_atom.unsqueeze(-1) + def forward( self, coord: torch.Tensor, @@ -100,8 +109,9 @@ def forward( ``mask_mag`` is built from the per-type spin gate on the local ``atype``; non-magnetic atoms already carry a zero magnetic force (the descriptor gates the spin embedding by type), so the force itself needs - no re-masking. This is the runtime counterpart of the static schema in - :meth:`translated_output_def`. + no re-masking. Phantom padding uses a lookup-safe type but is intersected + with the physical-atom mask, so it always reports ``False``. This is the + runtime counterpart of the static schema in :meth:`translated_output_def`. """ model_ret = self.forward_common( coord, @@ -113,14 +123,10 @@ def forward( charge_spin=charge_spin, spin=spin, ) - nf, nloc = atype.shape[:2] model_predict: dict[str, torch.Tensor] = { "atom_energy": model_ret["energy"], "energy": model_ret["energy_redu"], - "mask_mag": self.spin_mask.index_select(0, atype.reshape(-1)).reshape( - nf, nloc, 1 - ) - > 0.0, + "mask_mag": self._make_mask_mag(atype), } if self.do_grad_r("energy"): model_predict["force"] = rearrange( @@ -397,8 +403,9 @@ def _attach_spin_masks( Internal SeZM lower outputs; ``energy_derv_r_mag`` has the per-local-atom shape ``(nf, nloc, 1, 3)``. atype - Local atom types with shape ``(nf, nloc)``, used to build - ``mask_mag`` with shape ``(nf, nloc, 1)``. + Atom types aligned with the magnetic-force rows before optional + ghost padding, used to build ``mask_mag`` with a trailing singleton + axis. Negative phantom sentinels are allowed. nall Extended atom count the magnetic force is padded to. @@ -423,12 +430,7 @@ def _attach_spin_masks( nf, nloc = derv_r_mag.shape[:2] ghost_pad = derv_r_mag.new_zeros(nf, nall - nloc, *derv_r_mag.shape[2:]) model_ret["energy_derv_r_mag"] = torch.cat([derv_r_mag, ghost_pad], dim=1) - model_ret["mask_mag"] = ( - self.spin_mask.index_select(0, atype.reshape(-1)).reshape( - atype.shape[0], atype.shape[1], 1 - ) - > 0.0 - ) + model_ret["mask_mag"] = self._make_mask_mag(atype) return model_ret # ========================================================================= diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 5a190fe5af..1101127e1a 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -289,34 +289,22 @@ def get_data_loader( def get_dataloader_and_iter_lmdb( _data: LmdbDataset, ) -> tuple[LmdbBatchDataLoader, Generator[Any, None, None]]: - if _data.mixed_batch: - # TODO [mixed_batch=True]: Replace SameNlocBatchSampler with - # RandomSampler(replacement=False) + padding collate_fn. - # Changes needed: - # 1. _collate_lmdb_batch: pad coord/force/atype to max_nloc, - # add "atom_mask" bool tensor (nframes, max_nloc) - # 2. Use RandomSampler(_data, replacement=False) as sampler - # 3. Use fixed batch_size in DataLoader (not batch_sampler) - # 4. Model forward: apply atom_mask to descriptor/fitting - # 5. Loss: mask out padded atoms in force loss - raise NotImplementedError( - "mixed_batch=True training is not yet supported." - ) - # mixed_batch=False: group frames by nloc, each batch same nloc. - # SameNlocBatchSampler yields list[int] per batch, all same nloc. - # Auto batch_size is computed per-nloc-group inside the sampler. + # The sampler yields one list[int] per batch. Every batch is + # homogeneous in label availability; whether it is also + # homogeneous in atom count follows from the dataset's + # batch_size rule, which the sampler reads off the reader. from deepmd.dpmodel.utils.lmdb_data import ( - SameNlocBatchSampler, + LmdbBatchSampler, ) _block_targets = getattr(_data, "_block_targets", None) if self.world_size > 1: from deepmd.dpmodel.utils.lmdb_data import ( - DistributedSameNlocBatchSampler, + DistributedLmdbBatchSampler, ) - _inner_sampler = DistributedSameNlocBatchSampler( + _inner_sampler = DistributedLmdbBatchSampler( _data._reader, rank=self.rank, world_size=self.world_size, @@ -325,7 +313,7 @@ def get_dataloader_and_iter_lmdb( block_targets=_block_targets, ) else: - _inner_sampler = SameNlocBatchSampler( + _inner_sampler = LmdbBatchSampler( _data._reader, shuffle=True, block_targets=_block_targets, @@ -1595,7 +1583,11 @@ def log_loss_valid(_task_key: str = "Default") -> dict: task_key=_task_key, ) # more_loss.update({"rmse": math.sqrt(loss)}) - natoms = int(input_dict["atype"].shape[-1]) + # The metrics are per-atom quantities, so each batch + # weighs by the real atoms it holds summed over its + # frames. Phantom atoms (atype < 0), which pad a + # mixed-nloc batch, contribute to none of them. + natoms = int((input_dict["atype"] >= 0).sum()) sum_natoms += natoms for k, v in more_loss.items(): if "l2_" not in k: diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 3fed8282dd..888d08692f 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """PyTorch LMDB dataset — thin wrapper around framework-agnostic LmdbDataReader.""" +import functools import logging from collections.abc import ( Iterator, @@ -18,12 +19,14 @@ from deepmd.dpmodel.utils.lmdb_data import ( LmdbBatchIterator, + LmdbBatchSampler, LmdbDataReader, + LmdbDecodeConfig, LmdbTestData, - SameNlocBatchSampler, collate_lmdb_frames, compute_block_targets, is_lmdb, + resolve_per_atom_keys, ) from deepmd.env import ( get_lmdb_num_workers, @@ -44,7 +47,10 @@ ] -def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: +def _collate_lmdb_batch( + batch: list[dict[str, Any]], + config: LmdbDecodeConfig, +) -> dict[str, Any]: """Collate a list of frame dicts into a torch batch dict. Pre-converts per-frame numpy arrays to CPU torch tensors (zero-copy when @@ -53,19 +59,23 @@ def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: collate yields a torch dict (``sid`` becomes a torch tensor automatically via ``array_api_compat``). - All frames in the batch must have the same nloc (enforced by - SameNlocBatchSampler when mixed_batch=False). For mixed_batch=True, - raises NotImplementedError. - """ - if len(batch) > 1: - atypes = [d.get("atype") for d in batch if d.get("atype") is not None] - if atypes and any(len(a) != len(atypes[0]) for a in atypes): - raise NotImplementedError( - "mixed_batch collation (frames with different atom counts " - "in the same batch) is not yet supported. " - "Padding + mask in collate_fn needed." - ) + Frames of different atom counts are padded to the batch maximum; the + padded slots carry the phantom atom type. Frames must still agree on + label availability, which :class:`LmdbBatchSampler` guarantees. + Parameters + ---------- + batch : list[dict[str, Any]] + Decoded frames to collate. + config : LmdbDecodeConfig + Decoder state whose data requirements identify the per-atom fields. + + Returns + ------- + dict[str, Any] + One collated batch of CPU tensors. + """ + per_atom_keys = resolve_per_atom_keys(batch[0], config) with torch.device("cpu"): torch_frames: list[dict[str, Any]] = [] for f in batch: @@ -78,7 +88,7 @@ def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: else: tf[key] = torch.as_tensor(val) torch_frames.append(tf) - return collate_lmdb_frames(torch_frames) + return collate_lmdb_frames(torch_frames, per_atom_keys) def _lmdb_batch_to_torch( @@ -100,15 +110,15 @@ def _lmdb_batch_to_torch( return converted -class _SameNlocBatchSamplerTorch(Sampler): - """Torch Sampler adapter around the framework-agnostic SameNlocBatchSampler. +class _LmdbBatchSamplerTorch(Sampler): + """Torch Sampler adapter around the framework-agnostic LmdbBatchSampler. PyTorch DataLoader with batch_sampler expects a Sampler that yields - lists of indices. This wraps SameNlocBatchSampler (or - DistributedSameNlocBatchSampler) to satisfy that. + lists of indices. This wraps LmdbBatchSampler (or + DistributedLmdbBatchSampler) to satisfy that. """ - def __init__(self, inner: SameNlocBatchSampler) -> None: + def __init__(self, inner: LmdbBatchSampler) -> None: self._inner = inner def __iter__(self) -> Iterator[list[int]]: @@ -140,7 +150,7 @@ def __init__( num_workers: int | None = None, ) -> None: self.dataset = dataset - self.batch_sampler = _SameNlocBatchSamplerTorch(sampler) + self.batch_sampler = _LmdbBatchSamplerTorch(sampler) self.sampler = sampler self._pin_memory = pin_memory self._batch_iterator = LmdbBatchIterator( @@ -188,9 +198,11 @@ class LmdbDataset(Dataset): - ``"max:N"``: ``max(1, floor(N / nloc))`` per nloc group. - ``"filter:N"``: same per-nloc formula as ``"max:N"`` and drops every frame whose ``nloc > N`` from the dataset. - mixed_batch : bool - If True, allow different nloc in the same batch (future). - If False (default), use SameNlocBatchSampler. + - ``"mix:N"``: mixed-nloc batching to a padded-slot budget of ``N``; + frames of different atom counts share a batch and the shorter ones + are padded with phantom atoms. + auto_prob_style : str, optional + ``auto_prob`` string used to reweight the original systems. """ def __init__( @@ -198,21 +210,16 @@ def __init__( lmdb_path: str, type_map: list[str], batch_size: int | str = "auto", - mixed_batch: bool = False, auto_prob_style: str | None = None, ) -> None: - self._reader = LmdbDataReader( - lmdb_path, type_map, batch_size, mixed_batch=mixed_batch + self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) + self._collate = functools.partial( + _collate_lmdb_batch, config=self._reader.decode_config ) - if mixed_batch: - # Future: DataLoader with padding collate_fn - raise NotImplementedError( - "mixed_batch=True is not yet supported. " - "Requires padding + mask in collate_fn." - ) - - # Compute block_targets from auto_prob_style if provided + # Compute block_targets from auto_prob_style if provided. An empty + # result means the configured probabilities need no reweighting, which + # is the common case and worth no log line of its own. self._block_targets = None if auto_prob_style is not None and self._reader.frame_system_ids is not None: self._block_targets = compute_block_targets( @@ -220,26 +227,25 @@ def __init__( self._reader.nsystems, self._reader.system_nframes, ) - if self._block_targets is not None: + if self._block_targets: log.info( f"LMDB auto_prob: {len(self._block_targets)} blocks, " f"nsystems={self._reader.nsystems}" ) - # Same-nloc batching: use SameNlocBatchSampler - sampler = SameNlocBatchSampler( + sampler = LmdbBatchSampler( self._reader, shuffle=True, block_targets=self._block_targets, ) - self._batch_sampler = _SameNlocBatchSamplerTorch(sampler) + self._batch_sampler = _LmdbBatchSamplerTorch(sampler) with torch.device("cpu"): self._inner_dataloader = DataLoader( self, batch_sampler=self._batch_sampler, num_workers=0, - collate_fn=_collate_lmdb_batch, + collate_fn=self._collate, ) # Per-nloc and label-availability dataloaders for make_stat_input. @@ -264,7 +270,7 @@ def _rebuild_nloc_dataloaders(self) -> None: shuffle=False, num_workers=0, drop_last=False, - collate_fn=_collate_lmdb_batch, + collate_fn=self._collate, ) dataloaders.append(dl) self._nloc_dataloaders = dataloaders @@ -295,8 +301,9 @@ def nframes(self) -> int: return self._reader.nframes @property - def mixed_batch(self) -> bool: - return self._reader.mixed_batch + def mixed_nloc(self) -> bool: + """Whether one batch may span several atom counts.""" + return self._reader.mixed_nloc @property def mixed_type(self) -> bool: @@ -411,8 +418,6 @@ def set_noise(self, noise_settings: dict[str, Any]) -> None: @property def index(self) -> list[int]: """Number of batches per logical LMDB dataset.""" - if not self._block_targets: - return self._reader.index return [self.total_batch] @property diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index df5f3dc119..3c68251a7d 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -353,6 +353,21 @@ def _matrix_to_extended_inputs( slot_idx = edge_idx % max_neighbors zero_shift = torch.all(shift == 0, dim=1) + # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical + # site, but the geometric search still pairs them with real atoms near + # their placeholder coordinates. Edges touching one are neutralized here, + # ahead of the ``sum(sel)`` truncation applied by the caller, so that a + # phantom can never displace a genuine neighbor from a real atom's list. + # Rather than compacting the edge arrays, which would cost a device sync, + # such an edge is rewritten to carry the empty marker and routed through + # the direct branch below: it then writes -1 into its own matrix slot, + # which is what that slot already holds, and is excluded from the ghost + # materialization of Step 3. + atype_flat = atype.reshape(-1) + real_pair = (atype_flat[dst] >= 0) & (atype_flat[src] >= 0) + src_local = torch.where(real_pair, src_local, -1) + zero_shift = zero_shift | ~real_pair + # === Step 2. Direct neighbors keep their local extended indices === # Zero-shift neighbors already live in the leading local block of # `extended_coord`, so their DeePMD nlist value is simply `src_local`. diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 06f57176cd..6aa5f4add4 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -254,6 +254,30 @@ def forward( do_atomic_virial=do_atomic_virial, neighbor_list=neighbor_list, ) + return self._to_public_keys(model_ret, do_atomic_virial) + + def _to_public_keys( + self, + model_ret: dict[str, torch.Tensor], + do_atomic_virial: bool, + ) -> dict[str, torch.Tensor]: + """Rename one ``call_common`` result to the public energy-model keys. + + The renaming is a property of the output definition, not of the node + axis, so it serves the rectangular and the ragged entry alike. + + Parameters + ---------- + model_ret : dict[str, torch.Tensor] + A ``call_common`` result, in internal ``_`` keys. + do_atomic_virial : bool + Whether the per-atom virial was requested and should be carried. + + Returns + ------- + dict[str, torch.Tensor] + The same tensors under the public names. + """ model_predict = {} model_predict["atom_energy"] = model_ret["energy"] model_predict["energy"] = model_ret["energy_redu"] @@ -263,12 +287,68 @@ def forward( model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) if do_atomic_virial: model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] + for key in ("mask", "n_node"): + if key in model_ret: + model_predict[key] = model_ret[key] if self.atomic_output_def()["energy"].r_hessian: model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3) return model_predict + def forward_ragged( + self, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Evaluate the energy model over a batch whose node axis is flat. + + The counterpart of :meth:`forward` for frames held concatenated rather + than padded to a common atom count. Arguments and results share their + meaning with :meth:`call_common_ragged`, whose per-atom outputs keep + the flat axis; only the keys are the public ones. + + Parameters + ---------- + coord : torch.Tensor + Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + atype : torch.Tensor + Local atom types with shape ``(N,)``. + n_node : torch.Tensor + Atoms per frame with shape ``(nf,)``. + box : torch.Tensor or None, optional + Simulation cell with shape ``(nf, 3, 3)``. + fparam : torch.Tensor or None, optional + Frame parameters with shape ``(nf, ndf)``. + aparam : torch.Tensor or None, optional + Atomic parameters with shape ``(N, nda)``. + do_atomic_virial : bool, default: False + Whether to return per-atom virials. + charge_spin : torch.Tensor or None, optional + Frame-level charge and spin conditioning with shape ``(nf, 2)``. + + Returns + ------- + dict[str, torch.Tensor] + Public energy-model keys; per-atom entries have leading dimension + ``N`` and per-frame entries ``nf``. + """ + model_ret = self.call_common_ragged( + coord, + atype, + n_node, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) + return self._to_public_keys(model_ret, do_atomic_virial) + def forward_lower( self, extended_coord: torch.Tensor, @@ -605,6 +685,7 @@ def forward_lower_graph_exportable_with_comm( ---------- atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial As in :meth:`forward_lower_graph_exportable`. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost The 8 comm tensors (see ``_make_comm_sample_inputs`` in ``serialization.py``), packed into ``comm_dict`` inside the @@ -621,6 +702,7 @@ def forward_lower_graph_exportable_with_comm( back with synchronizing D2H reads (``4 * nlayers`` per MD step). The C++ ``run_model_graph_with_comm`` implements this placement. + n_local (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node count consumed IN-GRAPH by the owned-node energy mask (it @@ -630,6 +712,7 @@ def forward_lower_graph_exportable_with_comm( access). Carries the same value as the ``nlocal`` comm tensor; the two inputs exist precisely to separate the device-compute role from the host-MPI-control role. + **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index c6def8f136..00254e3efb 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -20,6 +20,10 @@ from deepmd.dpmodel.output_def import ( OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + compact_nodes, + expand_node_values, +) from deepmd.kernels.utils import ( cuda_infer_level, ) @@ -28,6 +32,8 @@ ) from deepmd.pt_expt.utils.graph_builder import ( build_neighbor_graph_for_method, + build_ragged_neighbor_graph, + resolve_neighbor_graph_method, ) from deepmd.pt_expt.utils.graph_csr import ( validate_graph_csr_for_export, @@ -710,6 +716,102 @@ def _resolve_graph_method( return getattr(self, "neighbor_graph_method", "dense") return None + def call_common_ragged( + self, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Model forward over a batch whose node axis is already flat. + + The rectangular :meth:`call_common` pads frames of unequal atom + count to a common width and unpads its output again. A caller that + holds the frames concatenated skips both: the node axis it passes + in is the one the graph lower works on, and the per-atom outputs + come back on it. + + Parameters + ---------- + coord : torch.Tensor + Local coordinates with shape ``(N, 3)``, frame-major over + ``n_node``. + atype : torch.Tensor + Local atom types with shape ``(N,)``. + n_node : torch.Tensor + Atoms per frame with shape ``(nf,)``. + box : torch.Tensor or None, optional + Simulation cell with shape ``(nf, 3, 3)``, or ``None`` for + non-periodic. + fparam : torch.Tensor or None, optional + Frame parameter with shape ``(nf, ndf)``. + aparam : torch.Tensor or None, optional + Atomic parameter with shape ``(N, nda)``. + do_atomic_virial : bool, default: False + Whether to compute the atomic virial. + charge_spin : torch.Tensor or None, optional + Frame-level charge/spin conditioning with shape ``(nf, 2)``. + + Returns + ------- + dict[str, torch.Tensor] + The standard model dict. Per-atom keys keep the flat ``(N, *)`` + axis; per-frame keys have leading dimension ``nf``. + + Raises + ------ + NotImplementedError + If the model has no graph lower to read a flat node axis with. + """ + if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): + raise NotImplementedError( + "a flat node axis requires a mixed_types descriptor with a " + "graph lower; this model reads a rectangular one, so its " + "batches must be padded to a common atom count" + ) + # The trainer resolves ``auto`` once and installs the concrete + # builder on the model. A model reached outside it has none, and + # resolving against its own device is what keeps that case from + # silently taking the CPU builder on a GPU. + method = getattr(self, "neighbor_graph_method", None) + if method is None: + method = resolve_neighbor_graph_method("auto", coord.device) + graph = build_ragged_neighbor_graph( + method, + coord, + atype, + n_node, + box, + self.get_rcut(), + getattr(self.atomic_model, "pair_excl", None), + ) + predict = self.forward_common_lower_graph( + atype, + graph.n_node, + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + destination_sorted=graph.destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + # The per-atom mask a rectangular batch carries is what tells the + # loss each frame's real atom count. Nothing here is padded, so the + # counts are stated outright instead. + predict["n_node"] = graph.n_node + return predict + def _call_common_graph( self, cc: torch.Tensor, @@ -787,10 +889,25 @@ def _call_common_graph( method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr ) nf, nloc = atype.shape[:2] - atype_flat = atype.reshape(nf * nloc) + n_padded = nf * nloc + atype_flat = atype.reshape(n_padded) + # A batch of unequal atom counts arrives padded to a common width + # with phantom atoms (atype < 0). The builders leave them out of + # every edge, so dropping them from the node axis costs nothing and + # spares the network from evaluating them. On a batch of uniform + # atom count the mask is all true and this is a renumbering by the + # identity. + ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_flat = atype_flat[node_index] # graph-lower ABI: aparam/spin are FLAT on the node axis, (N, nda)/(N, 3). - ap_flat = ap.reshape(nf * nloc, ap.shape[-1]) if ap is not None else None - spin_flat = spin.reshape(nf * nloc, 3) if spin is not None else None + ap_flat = ( + ap.reshape(n_padded, ap.shape[-1])[node_index] + if ap is not None + else None + ) + spin_flat = ( + spin.reshape(n_padded, 3)[node_index] if spin is not None else None + ) model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -810,11 +927,14 @@ def _call_common_graph( charge_spin=charge_spin, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom - # outputs (N = nf * nloc for a carry-all rectangular graph). - # Unravel to rectangular ``(nf, nloc, *)`` at the public I/O boundary - # so that callers receive the same shape as the dense ``call_common``. - N = nf * nloc - # public call_common always passes rectangular (nf,nloc) coord/atype (N == nf*nloc), so this unravel always applies; ragged graphs reach call_lower_graph/forward_common_lower_graph directly (no unravel) and stay flat (N,*). + # outputs over the real atoms. Scatter them back onto the padded + # width and unravel to rectangular ``(nf, nloc, *)`` at the public + # I/O boundary, so that callers receive the same shape as the dense + # ``call_common``. A phantom slot reads zero, which is what a + # masked-out atom contributed there before. + N = node_index.shape[0] + # Only the rectangular entry reaches this scatter; the ragged + # one keeps the flat axis its caller handed over. for k in list(model_predict.keys()): v = model_predict[k] # per-frame reduced keys (..._redu) keep their (nf, *) shape; only node-level (N,*) keys unravel — guards the nloc==1 case where N == nf. @@ -823,7 +943,9 @@ def _call_common_graph( and not k.endswith("_redu") and v.shape[:1] == torch.Size([N]) ): - model_predict[k] = v.reshape(nf, nloc, *v.shape[1:]) + model_predict[k] = expand_node_values( + v, node_index, n_padded + ).reshape(nf, nloc, *v.shape[1:]) # Graph-native Hessian (parallel to the dense ``forward_common_atomic`` # loop): differentiate the reduced output w.r.t. the LOCAL coords by # rebuilding the graph inside the wrapper. Added AFTER the unravel so diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f277452fcb..fb9a2393e0 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1026,6 +1026,76 @@ def __getattr__(self, name: str) -> Any: except AttributeError: return getattr(self.original_model, name) + def forward_ragged( + self, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Compiled forward over a batch whose node axis is already flat. + + The compiled lower works on that axis in either case -- its trace keeps + the frame count, the node count and the edge count as independent + symbols -- so a ragged batch simply skips the padding round trip the + rectangular :meth:`forward` performs around it. + + Parameters + ---------- + coord : torch.Tensor + Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + atype : torch.Tensor + Local atom types with shape ``(N,)``. + n_node : torch.Tensor + Atoms per frame with shape ``(nf,)``. + box : torch.Tensor or None, optional + Simulation cell, ``(nf, 3, 3)`` or ``(nf, 9)``. + fparam : torch.Tensor or None, optional + Frame parameters with shape ``(nf, ndf)``. + aparam : torch.Tensor or None, optional + Atomic parameters with shape ``(N, nda)``. + do_atomic_virial : bool, default: False + Whether to return per-atom virials. + charge_spin : torch.Tensor or None, optional + Frame-level charge and spin conditioning with shape ``(nf, 2)``. + + Returns + ------- + dict[str, torch.Tensor] + Public model keys; per-atom entries keep the flat axis. + + Raises + ------ + NotImplementedError + If the model reads a rectangular node axis, which cannot represent + frames of unequal atom count without padding. + """ + del do_atomic_virial + if self._graph_eligible is None: + self._graph_eligible = model_uses_graph_lower(self.original_model) + if not self._graph_eligible: + raise NotImplementedError( + "a flat node axis requires a model whose descriptor reads one; " + "this model compiles the dense (nlist) lower, whose batches " + "must be padded to a common atom count" + ) + return self._forward_graph( + coord, + atype, + box, + fparam, + aparam, + charge_spin, + int(n_node.shape[0]), + 0, + self.original_model.get_rcut(), + n_node=n_node, + ) + def forward( self, coord: torch.Tensor, @@ -1252,6 +1322,7 @@ def _forward_graph( nframes: int, nloc: int, rcut: float, + n_node: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Carry-all GRAPH forward -> compiled ``forward_common_lower_graph``. @@ -1262,20 +1333,38 @@ def _forward_graph( so no extended->local scatter is needed; only the flat ``(N, *)`` node keys are unravelled to ``(nf, nloc, *)`` at the I/O boundary. """ + from deepmd.dpmodel.utils.neighbor_graph import ( + compact_nodes, + expand_node_values, + ) from deepmd.pt_expt.utils.graph_builder import ( build_neighbor_graph_for_method, + build_ragged_neighbor_graph, ) _model = self.original_model - coord_3d = coord.detach().reshape(nframes, nloc, 3) + # A ragged batch already holds the node axis the lower works on; a + # rectangular one is unravelled to it, and its per-atom outputs are + # folded back at the end. + ragged = n_node is not None + n_padded = nframes * nloc + # The builders take the shape the layout hands over: a flat node axis, + # or the rectangular one a padded batch carries. + coord_3d = ( + coord.detach().reshape(-1, 3) + if ragged + else coord.detach().reshape(nframes, nloc, 3) + ) box_flat = box.detach().reshape(nframes, 9) if box is not None else None # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- like # every per-node tensor of the graph schema (the trace sample from # build_synthetic_graph_inputs is flat too, so the compiled lower's - # input spec expects it). - if aparam is not None: - aparam = aparam.reshape(nframes * nloc, -1) + # input spec expects it). A ragged batch already carries it that way; + # a rectangular one may fold the component axis into its frame rows, + # so its node count is what unravels it. + if aparam is not None and not ragged: + aparam = aparam.reshape(n_padded, -1) # Mirror the optional-input defaulting of the dense path / eager # call_common: a model configured with fparam / charge_spin substitutes @@ -1314,15 +1403,26 @@ def _forward_graph( # into edge_mask here so the compiled lower consumes a pre-excluded graph # (the lower no longer re-applies it), matching the eager path exactly. pair_excl = getattr(_model.atomic_model, "pair_excl", None) - ng = build_neighbor_graph_for_method( - getattr(_model, "neighbor_graph_method", "dense"), - coord_3d, - atype, - box_flat, - rcut, - pair_excl, - ) - atype_flat = atype.reshape(nframes * nloc) + method = getattr(_model, "neighbor_graph_method", "dense") + if ragged: + ng = build_ragged_neighbor_graph( + method, coord_3d, atype, n_node, box_flat, rcut, pair_excl + ) + atype_flat, node_index = atype, None + else: + ng = build_neighbor_graph_for_method( + method, coord_3d, atype, box_flat, rcut, pair_excl + ) + # A rectangular batch of unequal atom counts is padded to a common + # width with phantom atoms (atype < 0). The builders leave them out + # of every edge, so dropping them from the node axis costs nothing + # and spares the network from evaluating them. On a batch of + # uniform atom count this is a renumbering by the identity. + atype_flat = atype.reshape(n_padded) + ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_flat = atype_flat[node_index] + if aparam is not None: + aparam = aparam[node_index] # Lazy compile of the GRAPH lower (cached per structure key). if self.compiled_forward_lower is None: @@ -1386,9 +1486,14 @@ def _forward_graph( # The compiled graph lower emits PUBLIC keys on the FLAT node axis # (``atom_energy`` / ``force`` are (N, *); ``energy`` / ``virial`` are - # (nf, *)). Unravel the node-level keys to rectangular (nf, nloc, *) so - # callers receive the same shapes as the dense path. - N = nframes * nloc + # (nf, *)). A ragged caller reads that axis directly. A rectangular one + # has its node-level keys scattered back onto the padded width and + # unravelled to (nf, nloc, *), where a phantom slot reads zero -- what + # a masked-out atom contributed there before. + if ragged: + result["n_node"] = ng.n_node + return result + N = node_index.shape[0] # Node-level (per-atom, lead dim N) public keys emitted by the graph # lower; the remaining keys are frame-level (lead dim nf) and must NOT # be unravelled. Keying on the NAME rather than the ``N != nframes`` @@ -1403,7 +1508,9 @@ def _forward_graph( and val is not None and val.shape[:1] == torch.Size([N]) ): - out[key] = val.reshape(nframes, nloc, *val.shape[1:]) + out[key] = expand_node_values(val, node_index, n_padded).reshape( + nframes, nloc, *val.shape[1:] + ) else: out[key] = val return out @@ -1986,6 +2093,7 @@ def update_finetune_bias( self._configure_neighbor_graph_method( training_params.get("neighbor_graph_method", "auto") ) + self._configure_batch_layout(training_data, validation_data) # torch.compile ------------------------------------------------------- if self.enable_compile: @@ -2081,6 +2189,47 @@ def _raise_if_full_validation_unsupported( # torch.compile helpers # ------------------------------------------------------------------ + def _configure_batch_layout(self, *data_maps: Any) -> None: + """Ask each LMDB data system for the layout its own model can consume. + + A model whose descriptor reads a flat node axis takes the frames of a + batch concatenated, which spares it the padding that frames of unequal + atom count would otherwise need. Every other model reads an + ``(nf, nloc, ...)`` axis and needs them padded to a common width. Only + the trainer sees both sides, and it settles the question here, once, + before any batch is drawn. + + A graph lower is necessary but not sufficient: the model must also + expose an entry that takes the flat axis, which the composed models + (linear, ZBL bridging) do not. Native-spin models also stay rectangular: + their public output translation needs the spin-specific force and mask, + while this generic ragged entry translates energy-model outputs only. + Requiring these capabilities keeps each model on the layout it can read. + + Each task has its own data system and its own model, so the answer is + each task's own: a multi-task run pairing a graph model with a dense + one gives the first concatenated batches and the second padded ones. + + Parameters + ---------- + *data_maps : Any + The training and validation data systems, either bare or as the + per-task mappings a multi-task run builds. + """ + for task_key in self.model_keys: + model = self.models[task_key] + ragged = ( + not model.has_spin() + and model_uses_graph_lower(model) + and hasattr(model, "forward_ragged") + ) + for data_map in data_maps: + data = ( + data_map.get(task_key) if isinstance(data_map, dict) else data_map + ) + if hasattr(data, "use_ragged_batches"): + data.use_ragged_batches(ragged) + def _configure_neighbor_graph_method(self, requested: str) -> None: """Resolve and install the training graph builder on eligible models.""" graph_models = [ @@ -2880,7 +3029,10 @@ def evaluate_validation( label=val_label, task_key=task.key, ) - natoms = int(val_input["atype"].shape[-1]) + # The metrics are per-atom quantities, so each batch weighs by the + # real atoms it holds summed over its frames. Phantom atoms + # (atype < 0), which pad a mixed-nloc batch, contribute to none. + natoms = int((val_input["atype"] >= 0).sum()) sum_natoms += natoms for key, value in vmore.items(): if "l2_" not in key: diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index f59b707217..634f3d4045 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -148,6 +148,7 @@ def forward( task_key: str | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + n_node: torch.Tensor | None = None, ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None, dict | None]: if not self.multi_task: task_key = "Default" @@ -164,6 +165,7 @@ def forward( "fparam": fparam, "aparam": aparam, "charge_spin": charge_spin, + "n_node": n_node, } # ``spin`` (native or virtual-atom magnetic moment) is only accepted # by spin-capable model forward()s; mirrors @@ -182,6 +184,10 @@ def forward( if label is None: return model_pred, None, None + # The width a rectangular batch pads its frames to. A ragged batch has + # none, so this is its whole atom count and the loss reads each frame's + # own from ``model_pred["n_node"]``; the terms that cannot express + # themselves per frame refuse such a batch rather than read this. natoms = atype.shape[-1] loss, more_loss = self.loss[task_key]( cur_lr, @@ -219,8 +225,17 @@ def _forward_without_loss( task_key: str, input_dict: dict[str, Any], ) -> dict[str, torch.Tensor]: - """Return model predictions without constructing a loss.""" - return self.model[task_key](**input_dict) + """Return model predictions without constructing a loss. + + ``n_node`` marks a batch whose frames are concatenated rather than + padded to a common atom count. Its node axis is already the one the + graph lower works on, so it takes the entry that skips the padding + round trip; ``forward`` accepts only the rectangular shape. + """ + model = self.model[task_key] + if input_dict.get("n_node") is None: + return model(**{k: v for k, v in input_dict.items() if k != "n_node"}) + return model.forward_ragged(**input_dict) def set_extra_state(self, state: dict) -> None: self.model_params = state.get("model_params", {}) diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py index 4e916abea0..b2cbad736d 100644 --- a/deepmd/pt_expt/utils/edge_schema.py +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -18,6 +18,37 @@ _DUMMY_EDGE_COUNT = 2 +def real_atom_edge_mask( + atype_flat: torch.Tensor, + src: torch.Tensor, + dst: torch.Tensor, +) -> torch.Tensor: + """Keep only the edges whose two endpoints are both real atoms. + + Phantom atoms (``atype < 0``) pad a mixed-nloc batch to a rectangular + shape: they occupy a tensor slot but no physical site. A geometric + neighbor search cannot know that, and reports them as neighbors of + whichever real atoms lie near their placeholder coordinates. Dropping + those edges restores the invariant that a phantom neither carries an + environment of its own nor enters a real atom's. + + Parameters + ---------- + atype_flat : torch.Tensor + Atom types flattened over the batch, with shape ``(nf * nloc,)``. + src : torch.Tensor + Source endpoint of each edge, indexing ``atype_flat``. + dst : torch.Tensor + Destination endpoint of each edge, indexing ``atype_flat``. + + Returns + ------- + torch.Tensor + Boolean mask over edges, with shape ``(nedge,)``. + """ + return (atype_flat[src] >= 0) & (atype_flat[dst] >= 0) + + def _append_dummy_edges( edge_index: torch.Tensor, edge_vec: torch.Tensor, @@ -119,8 +150,21 @@ def edge_schema_from_extended( # No ``edge_len2 <= rcut**2`` upper bound here: ``nlist`` is contractually # cutoff-truncated by the caller (see the docstring). Only padding (-1), - # ghost-only neighbours, and coincident pairs are dropped. - edge_keep = valid_flat & (src_local >= 0) & (src_local < nloc) & (edge_len2 > 1e-10) + # ghost-only neighbours, coincident pairs and phantom endpoints are + # dropped. Without ``mapping`` the source index spans the extended atoms, + # so it is clamped before the type lookup: entries the preceding terms + # already reject must still not index out of bounds. + edge_keep = ( + valid_flat + & (src_local >= 0) + & (src_local < nloc) + & (edge_len2 > 1e-10) + & real_atom_edge_mask( + atype[:, :nloc].reshape(-1), + src_actual.clamp(0, nf * nloc - 1), + dst_actual, + ) + ) valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() edge_index = torch.stack( [ @@ -206,7 +250,11 @@ def edge_schema_from_neighbor_matrix( edge_vec_all.index_add_(0, shifted_idx, shift_cart) edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) - edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut)) + edge_keep = ( + (edge_len2 > 1e-10) + & (edge_len2 <= float(rcut) * float(rcut)) + & real_atom_edge_mask(atype.reshape(-1), src_actual, dst) + ) valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() schema = _append_dummy_edges( torch.stack( @@ -272,7 +320,11 @@ def edge_schema_from_ij_shifts( (sel_shifts[:, :, None] * cell).sum(1), ) edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) - edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut)) + edge_keep = ( + (edge_len2 > 1e-10) + & (edge_len2 <= float(rcut) * float(rcut)) + & real_atom_edge_mask(atype.reshape(-1), jj, ii) + ) valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() schema = _append_dummy_edges( torch.stack( diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index b074a3af8f..aa6944bfa3 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -8,6 +8,10 @@ import torch +from deepmd.dpmodel.utils.lmdb_data import ( + PHANTOM_ATOM_TYPE, +) + if TYPE_CHECKING: from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, @@ -82,6 +86,82 @@ def resolve_neighbor_graph_method( return "nv" +def build_ragged_neighbor_graph( + method: str, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + pair_excl: "PairExcludeMask | None", + *, + with_csr: bool = False, +) -> "NeighborGraph": + """Build a carry-all graph over a batch whose node axis is already flat. + + The searches all take a rectangular tensor -- ``dense`` compares every pair + of one, and the others derive their per-frame bounds from its shape -- so + the frames are widened to a common width here and the resulting graph is + narrowed back onto the flat axis. Widening is a scatter and narrowing a + renumbering that drops no edge, since no builder draws one to a padded + slot. ``nv`` additionally withholds the padded slots from the search + itself, so for that backend the widening does not reach the geometry. + + Parameters + ---------- + method : str + The concrete builder, as resolved by :func:`resolve_neighbor_graph_method`. + coord : torch.Tensor + Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + atype : torch.Tensor + Local atom types with shape ``(N,)``. + n_node : torch.Tensor + Atoms per frame with shape ``(nf,)``. + box : torch.Tensor or None + Simulation cell with shape ``(nf, 3, 3)``, or ``None`` for non-periodic. + rcut : float + Cutoff radius. + pair_excl : PairExcludeMask or None + Model-level pair exclusion, folded into the edge mask at build time. + with_csr : bool, default: False + Whether to attach destination/source CSR views. + + Returns + ------- + NeighborGraph + A graph whose node axis is the one described by ``n_node``. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + compact_nodes, + ) + + nf = int(n_node.shape[0]) + width = int(n_node.max()) if nf else 0 + # Position of each flat node in the padded batch, built without a pass over + # the frames: its frame times the width, plus its rank within that frame. + frame = torch.repeat_interleave( + torch.arange(nf, dtype=n_node.dtype, device=n_node.device), n_node + ) + offset = torch.cumsum(n_node, 0) - n_node + slot = ( + torch.arange(int(n_node.sum()), dtype=n_node.dtype, device=n_node.device) + - offset[frame] + ) + padded_index = frame * width + slot + + padded_coord = coord.new_zeros((nf * width, 3)) + padded_atype = atype.new_full((nf * width,), PHANTOM_ATOM_TYPE) + padded_coord[padded_index] = coord + padded_atype[padded_index] = atype + padded_coord = padded_coord.reshape(nf, width, 3) + padded_atype = padded_atype.reshape(nf, width) + + graph = build_neighbor_graph_for_method( + method, padded_coord, padded_atype, box, rcut, pair_excl, with_csr=with_csr + ) + return compact_nodes(graph, padded_atype.reshape(-1) >= 0)[0] + + def build_neighbor_graph_for_method( method: str, coord: torch.Tensor, diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 7608d9eb61..d645ff9955 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -14,10 +14,10 @@ ) from deepmd.dpmodel.utils.lmdb_data import ( - DistributedSameNlocBatchSampler, + DistributedLmdbBatchSampler, LmdbBatchIterator, + LmdbBatchSampler, LmdbDataReader, - SameNlocBatchSampler, collect_lmdb_sampling_groups, compute_block_targets, ) @@ -39,11 +39,11 @@ class LmdbDataSystem: ``get_nsystems()``, and the ``nbatches``/``sys_probs`` pair from which the trainer derives an epoch length. The whole LMDB counts as one logical system. Internally uses :class:`LmdbDataReader` for I/O and - :class:`SameNlocBatchSampler`, or its distributed wrapper, to draw - same-nloc batches. Statistics use a separate logical-system view in which - every ``(nloc, label-availability)`` group is sampled independently, - matching the training sampler without changing the identity of the LMDB as - one training dataset. + :class:`LmdbBatchSampler`, or its distributed wrapper, to draw batches. + Statistics use a separate logical-system view in which every + ``(nloc, label-availability)`` group is sampled independently, matching + the training sampler without changing the identity of the LMDB as one + training dataset. Parameters ---------- @@ -52,12 +52,13 @@ class LmdbDataSystem: type_map Global type map from the model config. batch_size - Batch size spec; ``int``, ``"auto"``, or ``"auto:N"``. + Batch size spec; ``int``, ``"auto"``, ``"auto:N"``, ``"max:N"``, + ``"filter:N"``, or ``"mix:N"`` for mixed-nloc batching. auto_prob_style Optional ``auto_prob`` string (e.g. ``"prob_sys_size"``) for per-system reweighting via :func:`compute_block_targets`. seed - Optional seed for the shuffle in :class:`SameNlocBatchSampler`. + Optional seed for the shuffle in :class:`LmdbBatchSampler`. num_workers Number of LMDB decoder worker processes. ``None`` selects the hardware-aware default; zero or one disables multiprocessing. @@ -65,7 +66,7 @@ class LmdbDataSystem: Rank of this process in distributed training. world_size Number of distributed training processes. Values greater than one - select :class:`DistributedSameNlocBatchSampler`. + select :class:`DistributedLmdbBatchSampler`. """ def __init__( @@ -79,9 +80,7 @@ def __init__( rank: int = 0, world_size: int = 1, ) -> None: - self._reader = LmdbDataReader( - lmdb_path, type_map, batch_size, mixed_batch=False - ) + self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) block_targets = None if auto_prob_style is not None and self._reader.frame_system_ids is not None: @@ -92,23 +91,23 @@ def __init__( ) if world_size > 1: - distributed_sampler = DistributedSameNlocBatchSampler( - self._reader, - rank=rank, - world_size=world_size, - shuffle=True, - seed=seed, - block_targets=block_targets, + self._sampler: LmdbBatchSampler | DistributedLmdbBatchSampler = ( + DistributedLmdbBatchSampler( + self._reader, + rank=rank, + world_size=world_size, + shuffle=True, + seed=seed, + block_targets=block_targets, + ) ) - self._sampler = distributed_sampler else: - sampler = SameNlocBatchSampler( + self._sampler = LmdbBatchSampler( self._reader, shuffle=True, seed=seed, block_targets=block_targets, ) - self._sampler = sampler self._refresh_stat_groups() num_workers = ( get_lmdb_num_workers() if num_workers is None else int(num_workers) @@ -128,6 +127,20 @@ def _refresh_stat_groups(self) -> None: # pt_expt trainer surface # ------------------------------------------------------------------ + def use_ragged_batches(self, ragged: bool) -> None: + """Select the layout :meth:`get_batch` delivers. + + See :meth:`deepmd.dpmodel.utils.lmdb_data.LmdbDataReader.use_ragged_batches`. + The trainer calls this once it knows whether the model reads a flat + node axis, before training draws its first batch. + + Parameters + ---------- + ragged : bool + Whether to concatenate frames instead of padding them. + """ + self._reader.use_ragged_batches(ragged) + def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: """Return one batch as a numpy dict. @@ -154,6 +167,12 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: ------ IndexError If ``sys_idx`` does not identify an available statistical group. + + Notes + ----- + The batch is rectangular whatever layout training uses: output + statistics accumulate over an ``(nf, nloc, ...)`` axis. A group is + uniform in atom count, so that layout pads nothing. """ if not 0 <= sys_idx < len(self._stat_groups): raise IndexError( @@ -168,7 +187,7 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: start = 0 stop = min(start + batch_size, len(group_indices)) self._stat_offsets[sys_idx] = stop - return self._reader.decode_batch(group_indices[start:stop]) + return self._reader.decode_batch(group_indices[start:stop], ragged=False) def get_stat_nsystems(self) -> int: """Return the number of homogeneous statistical systems.""" @@ -189,13 +208,12 @@ def get_stat_numb_batches(self, sys_idx: int) -> int: def add_data_requirements( self, data_requirement: list[DataRequirementItem] ) -> None: - # Batches are partitioned by label availability. The sampler derives - # the partition on its first draw; only the distributed batch count is - # cached, so it is refreshed after the requirements change. + # Batches are partitioned by label availability, so new requirements + # repartition the frames. Both the statistical groups and the pass the + # sampler holds pending are therefore rebuilt from the new partition. self._reader.add_data_requirement(data_requirement) self._refresh_stat_groups() - if isinstance(self._sampler, DistributedSameNlocBatchSampler): - self._sampler.refresh_batch_count() + self._sampler.refresh_batch_count() def close(self) -> None: """Cancel prefetched work and release decoder processes.""" @@ -217,9 +235,7 @@ def get_nsystems(self) -> int: @property def nbatches(self) -> list[int]: """Return the global batch count of one full pass.""" - if isinstance(self._sampler, DistributedSameNlocBatchSampler): - return [self._sampler.total_batches] - return [len(self._sampler)] + return [self._sampler.total_batches] @property def sys_probs(self) -> list[float]: diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index e78fea3e16..aacd1383bc 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -60,6 +60,7 @@ def nv_matrix_to_ijs( num_neighbors: torch.Tensor, shifts: torch.Tensor, nloc: int, + node_index: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Decode nvalchemiops' dense neighbor matrix to a sparse edge list. @@ -72,25 +73,30 @@ def nv_matrix_to_ijs( ---------- neighbor_matrix (total_atoms, max_neighbors) int; ``neighbor_matrix[dst, slot] = src``, - both flattened batch indices in ``[0, total_atoms)``. Frames are + both flattened search indices in ``[0, total_atoms)``. Frames are batch-isolated: a neighbor always shares its center's frame. num_neighbors (total_atoms,) int, valid slot count per center. shifts (total_atoms, max_neighbors, 3) int periodic image shifts per slot. nloc - Atoms per frame (``total_atoms = nf * nloc``). + Atoms per frame of the rectangular batch the indices are reported on. + node_index + (total_atoms,) int, position of each searched atom in that rectangular + batch. Given when the search ran over a subset of the batch, such as + its real atoms alone; ``None`` when it ran over every slot, in which + case the search index and the batch position coincide. Returns ------- center_local - (E,) int64 per-frame local center index ``i`` (``dst % nloc``). + (E,) int64 per-frame local center index ``i``. src_local - (E,) int64 per-frame local neighbor index ``j`` (``src % nloc``). + (E,) int64 per-frame local neighbor index ``j``. shift (E, 3) int64 periodic image shift ``S``. frame_idx - (E,) int64 frame of each edge (``dst // nloc``). + (E,) int64 frame of each edge. """ device = neighbor_matrix.device total_atoms, max_neighbors = neighbor_matrix.shape @@ -103,6 +109,12 @@ def nv_matrix_to_ijs( dst = edge_idx // max_neighbors # flattened center src = neighbor_matrix.reshape(-1).index_select(0, edge_idx).to(torch.int64) shift = shifts.reshape(-1, 3).index_select(0, edge_idx).to(torch.int64) + if node_index is not None: + # Lift both endpoints from the search axis back onto the rectangular + # batch, so that the frame and local indices below are the ones every + # other builder reports and the caller needs no special case. + dst = node_index.index_select(0, dst) + src = node_index.index_select(0, src) frame_idx = (dst // nloc).to(torch.int64) # frame of the edge center_local = (dst % nloc).to(torch.int64) # i = center src_local = (src % nloc).to(torch.int64) # j = neighbor @@ -114,6 +126,7 @@ def nv_search_matrix( box: torch.Tensor | None, rcut: float, start_capacity: int, + node_index: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -136,12 +149,19 @@ def nv_search_matrix( rcut : cutoff radius. start_capacity : initial max-neighbor capacity; grown automatically when any atom has more neighbors than the current capacity. + node_index : (total_atoms,) positions, in the flattened batch, of the atoms + to search over, frame-major and ascending. ``None`` searches every slot. + The search itself is ragged -- nvalchemiops takes a flat position array + with explicit per-frame bounds -- so restricting it to the real atoms of + a padded batch keeps the phantom slots out of a cost that grows with the + square of the frame width. Returns ------- coord : (nf, nloc, 3) coordinates, normalized in-cell if periodic. cell : (nf, 3, 3) float box, or ``None`` for non-periodic. - neighbor_matrix : (total_atoms, capacity) int neighbor matrix. + neighbor_matrix : (total_atoms, capacity) int neighbor matrix, indexed on + the searched atoms rather than on the batch slots. num_neighbors : (total_atoms,) valid neighbor count per center. shifts : (total_atoms, capacity, 3) int periodic image shifts. """ @@ -163,16 +183,24 @@ def nv_search_matrix( cell = None pbc = None - total_atoms = nf * nloc - positions = coord.reshape(total_atoms, 3).detach() - batch_idx = torch.arange( - nf, dtype=torch.int32, device=device - ).repeat_interleave(nloc) - batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc - method = choose_nv_nlist_method(nloc, periodic=periodic, device=device) + positions = coord.reshape(nf * nloc, 3).detach() + if node_index is None: + batch_idx = torch.arange( + nf, dtype=torch.int32, device=device + ).repeat_interleave(nloc) + batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc + widest_frame = nloc + else: + positions = positions.index_select(0, node_index) + batch_idx = (node_index // nloc).to(torch.int32) + counts = torch.bincount(batch_idx.to(torch.int64), minlength=nf) + batch_ptr = torch.zeros(nf + 1, dtype=torch.int32, device=device) + batch_ptr[1:] = torch.cumsum(counts, 0).to(torch.int32) + widest_frame = int(counts.max()) if nf > 0 else 0 + method = choose_nv_nlist_method(widest_frame, periodic=periodic, device=device) extra_nl_kwargs: dict[str, Any] = {} if method == "batch_naive": - extra_nl_kwargs["max_atoms_per_system"] = int(nloc) + extra_nl_kwargs["max_atoms_per_system"] = int(widest_frame) search_capacity = start_capacity while True: @@ -283,7 +311,7 @@ def build_neighbor_graph_nv( coord, box, empty_i, - nloc, + torch.full((nf,), nloc, dtype=torch.int64, device=device), layout=layout, with_csr=with_csr, canonicalize=canonicalize, @@ -300,25 +328,24 @@ def build_neighbor_graph_nv( 64, estimate_max_neighbors(float(rcut), atomic_density=0.25), ) + # Virtual atoms (atype < 0) are excluded as centers AND neighbors -- the + # World-2 builder contract shared with the dense reference builder. Here + # they are withheld from the search rather than filtered out of its result: + # the phantom slots of a mixed-nloc batch would otherwise widen every frame + # the search sees, and its cost grows with the square of that width. + atype_flat = torch.as_tensor(atype, device=device).reshape(nf * nloc) + node_index = torch.nonzero(atype_flat >= 0, as_tuple=False).flatten() coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix( - coord, box, rcut, start_capacity=initial_capacity + coord, box, rcut, start_capacity=initial_capacity, node_index=node_index ) box_out = cell # edge_vec is recomputed from these (normalized) coords - # Decode the dense matrix to a sparse (i, j, S) edge list (CPU-testable - # helper; see nv_matrix_to_ijs). + # Decode the dense matrix to a sparse (i, j, S) edge list, lifted back onto + # the rectangular batch (CPU-testable helper; see nv_matrix_to_ijs). center_local, src_local, shift, frame_idx = nv_matrix_to_ijs( - neighbor_matrix, num_neighbors, shifts, nloc + neighbor_matrix, num_neighbors, shifts, nloc, node_index=node_index ) - # virtual atoms (atype < 0) are excluded as centers AND neighbors — the - # World-2 builder contract shared with the dense reference builder; the - # geometric search above cannot know about them. - at = torch.as_tensor(atype, device=device).reshape(nf, nloc) - keep = (at[frame_idx, center_local] >= 0) & (at[frame_idx, src_local] >= 0) - center_local, src_local = center_local[keep], src_local[keep] - shift, frame_idx = shift[keep], frame_idx[keep] - graph = neighbor_graph_from_ijs( center_local, src_local, @@ -326,12 +353,11 @@ def build_neighbor_graph_nv( coord, box_out, frame_idx, - nloc, + torch.full((nf,), nloc, dtype=torch.int64, device=device), layout=layout, ) if pair_excl is not None: - at_flat = torch.as_tensor(atype, device=device).reshape(-1) - graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) if with_csr or canonicalize: graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) return graph diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index a715189ae5..2fce680896 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -165,7 +165,7 @@ def build_neighbor_graph_vesin( coord, box, empty_nf, - nloc, + torch.full((nf,), nloc, dtype=torch.int64, device=dev), layout=layout, with_csr=with_csr, canonicalize=canonicalize, @@ -222,7 +222,7 @@ def build_neighbor_graph_vesin( coord, box, nf_all, - nloc, + torch.full((nf,), nloc, dtype=torch.int64, device=dev), layout=layout, ) if pair_excl is not None: diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 2d28e1b090..7453176d32 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -231,6 +231,14 @@ def _build_single( ii, jj, ss = vesin_search_ijs( positions.detach(), cell if periodic else None, periodic, rcut, device ) + # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical + # site; the geometric search above cannot know that. Their pairs are + # dropped before the ``sum(sel)`` truncation below, so that a phantom can + # never displace a genuine neighbor from a real atom's list. Compacting + # the pair list, rather than marking the pairs empty, also keeps them out + # of the ``max_nn`` that sizes the dense candidate matrix. + real_pair = (atype[ii] >= 0) & (atype[jj] >= 0) + ii, jj, ss = ii[real_pair], jj[real_pair], ss[real_pair] # ss is int64 from the helper; cast to the coordinate dtype for the image sum. ss = ss.to(positions.dtype) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index eae723bc3d..bef85f5fae 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5169,6 +5169,7 @@ def training_data_args() -> list[ - string "mixed:N": the batch data will be sampled from all systems and merged into a mixed system with the batch size N. Only support the se_atten descriptor for TensorFlow backend.\n\n\ - string "max:N": automatically determines the batch size so that `batch_size * natoms` is at most `N`. `natoms` is the per-system atom count for npy data and the per-frame nloc for LMDB data. When a single system/frame already has more than `N` atoms, the batch size clamps to 1 and that batch will exceed `N`.\n\n\ - string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\ +- string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model and needs no configuration of its own: a descriptor reading a flat node axis (the graph route of the PyTorch Exportable backend) takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms with `atype = -1` that the neighbor list, the model and the loss all skip. Unlike `"max:N"`, which leaves an under-filled batch whenever an nloc group is small, this keeps every batch close to `N` atoms. A lone frame with nloc > `N` still forms a batch of its own.\n\n\ If MPI is used, the value should be considered as the batch size per task.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ @@ -5264,7 +5265,8 @@ def validation_data_args() -> list[ - string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.\n\n\ - string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.\n\n\ - string "max:N": automatically determines the batch size so that `batch_size * natoms` is at most `N`. `natoms` is the per-system atom count for npy data and the per-frame nloc for LMDB data. When a single system/frame already has more than `N` atoms, the batch size clamps to 1 and that batch will exceed `N`.\n\n\ -- string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.' +- string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\ +- string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model: a descriptor reading a flat node axis takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms that the neighbor list, the model and the loss all skip.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index dfc4076a9f..cea270bd37 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -197,6 +197,12 @@ def __init__( bsi = 1 bs.append(bsi) self.batch_size = bs + elif words[0] == "mix": + raise RuntimeError( + "the 'mix' batch_size rule packs frames of unequal atom " + "count into one batch and is only available for LMDB " + "datasets on the pt and pt_expt backends" + ) else: raise RuntimeError("unknown batch_size rule " + words[0]) elif isinstance(self.batch_size, list): diff --git a/doc/data/system.md b/doc/data/system.md index bfadb8887f..f11057b55a 100644 --- a/doc/data/system.md +++ b/doc/data/system.md @@ -82,4 +82,8 @@ With these edited files, one can put together frames with the same `Natoms`, ins To put frames with different `Natoms` into the same system, one can pad systems by adding virtual atoms whose type is `-1`. Virtual atoms do not contribute to any fitting property, so the atomic property of virtual atoms (e.g. forces) should be given zero. +:::{note} +Per-atom loss terms (force, atomic energy, dos, tensor) sum over a batch's real labels and divide by their total count, so a frame counts in proportion to its real atom number. This differs from averaging each frame's own per-label mean only when the frames of one batch pad a *different* number of virtual atoms; where the padding is uniform — including the case of no padding at all — the two are identical. Writing `real_atom_types.npy` with a per-frame varying number of `-1` is therefore the one configuration whose loss values, and hence training trajectory, differ from a per-frame average. This applies to the PyTorch backend; the TensorFlow backend normalizes by the padded width regardless. +::: + The API to generate or transfer to `mixed_type` format is available on [dpdata](https://github.com/deepmodeling/dpdata) for a more convenient experience. diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index e756f63b6c..84cfd30f1a 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -81,11 +81,26 @@ The sections {ref}`training_data ` and {ref}`validation_ - The key {ref}`batch_size ` specifies the number of frames used to train or validate the model in a training step. It can be set to - `list`: the length of which is the same as the {ref}`systems`. The batch size of each system is given by the elements of the list. + - `int`: all systems use the same batch size. + - `"auto"`: the same as `"auto:32"`, see `"auto:N"` + - `"auto:N"`: automatically determines the batch size so that the {ref}`batch_size ` times the number of atoms in the system is **no less than** `N`. + - `"max:N"`: automatically determines the batch size so that the {ref}`batch_size ` times the number of atoms in the system is **no more than** `N`. The minimum batch size is 1. **Supported backends**: PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }} + - `"filter:N"`: the same as `"max:N"` but removes the systems with the number of atoms larger than `N` from the data set. Throws an error if no system is left in a dataset. **Supported backends**: PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }} + + - `"mix:N"`: mixed-nloc batching for LMDB data sets. Frames of different atom counts share a batch, which is closed only when the next frame would push its atom axis past `N`. How that axis is laid out follows from the model and needs no configuration of its own: + + - A descriptor reading a **flat** node axis — the graph route of the PyTorch Exportable backend — takes the frames of a batch concatenated. `N` counts real atoms, nothing is padded, and the frames are packed in the shuffled order they arrive in, so one batch stays mixed in system size. + + - Every other descriptor reads a **rectangular** `(nframes, nloc, ...)` axis, on which the frames of a batch must share one atom count. Frames are therefore sorted by atom count and padded up to the widest one of their batch, which makes `N` the padded-slot count `nframes * max_nloc`. The padded slots hold phantom atoms marked `atype = -1`, which the neighbor list, the model and the loss all skip. Sorting is what keeps their number small, and it also makes the packing optimal: no arrangement of the same frames under the same budget yields fewer batches. + + Either way, a lone frame with more than `N` atoms still forms a batch of its own. Every other rule keeps a batch uniform in atom count, which leaves batches under-filled wherever an atom-count group is small; `"mix:N"` keeps them all close to `N` atoms instead. **Supported backends**: PyTorch {{ pytorch_icon }} (rectangular layout only), PyTorch Exportable + + Filling batches also makes the *sampling weight* more faithful. Per-atom loss terms (force, atomic energy) pool over a batch's real labels, so a frame already counts in proportion to its atom count there. Frame-level terms (energy, virial) weigh the frames of a batch equally, so a frame counts for `1 / nframes`, and an atom budget makes `nframes` follow the atom count. `"max:N"` distorts that wherever an atom-count group is too small to fill a batch — its few frames form a short batch and each is over-weighted many times over — while `"mix:N"` packs them with their neighbours and lands much closer to the intended weighting. - The key {ref}`numb_batch ` in {ref}`validate_data ` gives the number of batches of model validation. Note that the batches may not be from the same system The section {ref}`mixed_precision ` specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below: diff --git a/source/tests/common/dpmodel/test_from_ijs.py b/source/tests/common/dpmodel/test_from_ijs.py index adce99c5e5..a3dcad7e58 100644 --- a/source/tests/common/dpmodel/test_from_ijs.py +++ b/source/tests/common/dpmodel/test_from_ijs.py @@ -18,7 +18,7 @@ def test_edge_vec_and_index(self) -> None: j = np.array([1, 0]) # neighbor S = np.array([[0, 0, 0], [0, 0, 0]], dtype=np.int64) ng = neighbor_graph_from_ijs( - i, j, S, coord, box, nframe_id=np.zeros(2, np.int64), nloc=3 + i, j, S, coord, box, nframe_id=np.zeros(2, np.int64), n_node=np.array([3]) ) np.testing.assert_array_equal(ng.edge_index[0][ng.edge_mask], j) # src np.testing.assert_array_equal(ng.edge_index[1][ng.edge_mask], i) # dst @@ -34,7 +34,7 @@ def test_periodic_shift_in_edge_vec(self) -> None: j = np.array([1]) S = np.array([[-1, 0, 0]], dtype=np.int64) ng = neighbor_graph_from_ijs( - i, j, S, coord, box, nframe_id=np.zeros(1, np.int64), nloc=2 + i, j, S, coord, box, nframe_id=np.zeros(1, np.int64), n_node=np.array([2]) ) # coord[1] + (-1,0,0)@box - coord[0] = 5.5 - 6 - 0.5 = -1.0 np.testing.assert_allclose( diff --git a/source/tests/common/dpmodel/test_graph_ragged.py b/source/tests/common/dpmodel/test_graph_ragged.py index a651d245ea..6b1faae6c9 100644 --- a/source/tests/common/dpmodel/test_graph_ragged.py +++ b/source/tests/common/dpmodel/test_graph_ragged.py @@ -81,6 +81,66 @@ def test_forward_common_atomic_graph_ragged(): assert np.all(np.isfinite(out["energy"])) +def test_compact_nodes_recounts_frames_and_renumbers_edges(): + """A padded graph compacts onto its real nodes, frame blocks intact. + + Two frames padded to a width of four: frame 0 holds three real atoms and + frame 1 holds two, so the flat axis carries phantoms at 3, 6 and 7. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + compact_nodes, + expand_node_values, + ) + + node_mask = np.array([1, 1, 1, 0, 1, 1, 0, 0], dtype=bool) + graph = NeighborGraph( + n_node=np.array([4, 4], dtype=np.int64), + edge_index=np.array([[0, 1, 4], [1, 2, 5]], dtype=np.int64), + edge_vec=np.arange(9, dtype=np.float64).reshape(3, 3), + edge_mask=np.ones(3, dtype=bool), + ) + compacted, keep_index = compact_nodes(graph, node_mask) + + np.testing.assert_array_equal(compacted.n_node, [3, 2]) + np.testing.assert_array_equal(keep_index, [0, 1, 2, 4, 5]) + # Frame 1's nodes 4 and 5 sit at 3 and 4 once frame 0's phantom is gone. + np.testing.assert_array_equal(compacted.edge_index, [[0, 1, 3], [1, 2, 4]]) + np.testing.assert_array_equal( + frame_id_from_n_node(compacted.n_node), [0, 0, 0, 1, 1] + ) + # The edge payload is untouched: no edge is dropped, only renumbered. + np.testing.assert_array_equal(compacted.edge_vec, graph.edge_vec) + np.testing.assert_array_equal(compacted.edge_mask, graph.edge_mask) + + # Gathering onto the compact axis and expanding back is the identity on + # the real nodes and zero on the phantoms. + values = np.arange(8, dtype=np.float64)[:, None] + restored = expand_node_values(values[keep_index], keep_index, 8) + np.testing.assert_array_equal( + restored.ravel(), [0.0, 1.0, 2.0, 0.0, 4.0, 5.0, 0.0, 0.0] + ) + + +def test_compact_nodes_refuses_to_drop_a_node_that_carries_an_edge(): + """Renumbering a node with an edge would silently redirect that edge.""" + import pytest + + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + compact_nodes, + ) + + graph = NeighborGraph( + n_node=np.array([3], dtype=np.int64), + edge_index=np.array([[0], [1]], dtype=np.int64), + edge_vec=np.zeros((1, 3), dtype=np.float64), + edge_mask=np.ones(1, dtype=bool), + ) + with pytest.raises(ValueError, match="still carries an edge"): + compact_nodes(graph, np.array([1, 0, 1], dtype=bool)) + + def test_frame_id_rectangular(): fid = frame_id_from_n_node(np.array([4, 4], dtype=np.int64)) np.testing.assert_array_equal( diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 939d2b8a70..3537c97d15 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Unit tests for LmdbDataReader, LmdbTestData, SameNlocBatchSampler, etc. +"""Unit tests for LmdbDataReader, LmdbTestData, LmdbBatchSampler, etc. Pure dpmodel (NumPy/lmdb) tests — no PyTorch dependency. """ @@ -14,6 +14,9 @@ from concurrent.futures import ( Future, ) +from itertools import ( + pairwise, +) from pathlib import ( Path, ) @@ -30,12 +33,14 @@ from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module from deepmd.dpmodel.utils.lmdb_data import ( + DistributedLmdbBatchSampler, LmdbBatchIterator, + LmdbBatchSampler, LmdbDataReader, LmdbDecodeConfig, LmdbTestData, LmdbTestDataNlocView, - SameNlocBatchSampler, + _chop_mixed_nloc, _expand_indices_by_blocks, _merge_lmdb_chunks, _remap_atom_types, @@ -161,6 +166,46 @@ def _create_mixed_nloc_lmdb(path: str) -> str: return path +def _create_mix_probe_lmdb(path: str) -> str: + """Create an LMDB whose frame-level fields collide in shape with the atom axis. + + Frames alternate between 2 and 9 atoms and carry a two-component + ``fparam`` alongside a nine-component ``virials``. A padding rule that + guessed the atom axis from a leading dimension would misclassify + ``fparam`` on the 2-atom frames and ``virials`` on the 9-atom ones, so this + fixture pins the classification down to the exact ambiguous shapes. + """ + nlocs = [2, 9, 2, 9, 2, 9] + rng = np.random.RandomState(7) + env = lmdb.open(path, map_size=10 * 1024 * 1024) + with env.begin(write=True) as txn: + meta = { + "nframes": len(nlocs), + "frame_idx_fmt": "012d", + "frame_nlocs": nlocs, + "system_info": {"natoms": [1, 1], "formula": "probe"}, + } + txn.put(b"__metadata__", msgpack.packb(meta, use_bin_type=True)) + for idx, natoms in enumerate(nlocs): + frame = _make_frame(natoms=natoms, seed=idx) + frame["virials"] = { + "type": "= 0)) + np.testing.assert_array_equal(batch["coord"][0, 6:], 0.0) + np.testing.assert_array_equal(batch["force"][0, 6:], 0.0) + np.testing.assert_array_equal(batch["atype"][1], reader[8]["atype"]) + + # Atom counts stay the real per-frame values, not the padded width. + self.assertEqual(batch["natoms"][0, 0], 6) + self.assertEqual(batch["natoms"][1, 0], 12) + + def test_mix_pads_repeated_per_atom_fields(self): + """A ``repeat != 1`` requirement is stored flat but still padded by atom. + + ``atom_pref`` spends ``repeat`` leading entries per atom instead of + one, so its padded width is ``pad_nloc * repeat``. + """ + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36") + reader.add_data_requirement( + [DataRequirementItem("atom_pref", 1, atomic=True, must=False, repeat=3)] + ) + # Frames 0 and 8 hold 6 and 12 atoms respectively. + batch = reader.decode_batch([0, 8]) + self.assertEqual(batch["atom_pref"].shape, (2, 12 * 3)) + np.testing.assert_array_equal(batch["atom_pref"][0, 6 * 3 :], 0.0) + + def test_mix_chunked_decode_matches_in_process_decode(self): + """Splitting a padded batch across workers reproduces it exactly. + + Each chunk pads and lays out its fields independently, so a mixed-nloc + batch is the case where the chunks can disagree: they start on frames + of different atom counts. + """ + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36") + reader.add_data_requirement( + [ + DataRequirementItem("force", 3, atomic=True, must=False), + DataRequirementItem("atom_pref", 1, atomic=True, must=False, repeat=3), + ] + ) + # Frames 0 and 8 hold 6 and 12 atoms; one per chunk puts the wide + # frame second, so a per-chunk layout would size the chunks differently. + indices = [0, 8] + layout = reader.batch_layout(indices) + chunks = [ + lmdb_data_module.decode_lmdb_batch( + reader._transaction(), + [key], + reader.frame_format, + reader._decode_config, + layout.chunk(position, position + 1), + ) + for position, key in enumerate(reader.original_keys(indices)) + ] + merged = lmdb_data_module._merge_lmdb_chunks(chunks) + expected = reader.decode_batch(indices) + self.assertEqual(sorted(merged), sorted(expected)) + for field, value in expected.items(): + np.testing.assert_array_equal(merged[field], value, err_msg=field) + + def test_mix_keeps_frame_fields_of_ambiguous_shape(self): + """``fparam`` and ``virial`` are never padded, whatever the atom count.""" + path = _create_mix_probe_lmdb(f"{self._tmpdir.name}/probe.lmdb") + reader = LmdbDataReader(path, self._type_map, batch_size="mix:36") + reader.add_data_requirement( + [ + DataRequirementItem("virial", 9, atomic=False, must=False), + DataRequirementItem("fparam", 2, atomic=False, must=False), + DataRequirementItem("force", 3, atomic=True, must=False), + ] + ) + # A 2-atom frame makes fparam (2,) ambiguous; a 9-atom one does the + # same for virial (9,). Batching them together exercises both. + batch = reader.decode_batch([0, 1]) + self.assertEqual(batch["fparam"].shape, (2, 2)) + self.assertEqual(batch["virial"].shape, (2, 9)) + self.assertEqual(batch["coord"].shape, (2, 9, 3)) + np.testing.assert_array_equal(batch["fparam"][0], reader[0]["fparam"]) + np.testing.assert_array_equal(batch["virial"][0], reader[0]["virial"]) + + def test_mix_tracks_the_atom_proportional_weighting(self): + """Packing to a budget must not distort how much each frame counts. + + A batch is one optimizer step, and its frame-level terms (energy, + virial) average over the frames present, so a frame's weight in them + is ``1 / k_b``. An atom budget asks for ``k_b ~ budget / nloc``, that + is a weight proportional to atom count, matching what the pooled + per-atom terms give those frames unconditionally. + Same-nloc batching misses that ideal wherever an atom-count group is + too small to fill a batch: its few frames form a short batch and each + of them is over-weighted many times over. Filling batches across atom + counts removes that failure mode, and the residual error is the + padding, so both goals improve together. + """ + path = f"{self._tmpdir.name}/skewed.lmdb" + # A few dominant atom counts plus a tail of groups too small to fill a + # batch, down to one holding a single frame. This is the shape of a + # real merged dataset in miniature, and the tail is where same-nloc + # batching goes wrong. + _create_mixed_sid_nloc_lmdb( + path, + system_specs=[ + (120, 6), + (80, 8), + (60, 10), + (3, 12), + (2, 14), + (2, 16), + (1, 20), + ], + type_map=self._type_map, + ) + budget = 120 + + def weighting_error(spec): + """Log-spread of the frame weight against the ideal, and its worst case.""" + reader = LmdbDataReader(path, self._type_map, batch_size=spec) + nlocs = np.asarray(reader.frame_nlocs, dtype=np.float64) + batches = LmdbBatchSampler(reader, shuffle=True, seed=0).batches() + weight = np.empty(len(nlocs)) + for batch in batches: + weight[batch] = 1.0 / len(batch) + slots = sum(len(b) * reader.batch_pad_nloc(b) for b in batches) + # Weights matter only up to a global scale, which is the learning + # rate, so normalize both sides to mean 1 before comparing. + ratio = (weight / weight.mean()) / (nlocs / nlocs.mean()) + return np.std(np.log(ratio)), ratio.max(), nlocs.sum() / slots + + same_spread, same_worst, same_efficiency = weighting_error(f"max:{budget}") + mix_spread, mix_worst, mix_efficiency = weighting_error(f"mix:{budget}") + + self.assertEqual(same_efficiency, 1.0) + self.assertGreater(mix_efficiency, 0.95) + # Both the typical and the worst-case deviation shrink by a wide + # margin; the thresholds leave room for the packing to change. + self.assertLess(mix_spread, 0.5 * same_spread) + self.assertLess(mix_worst, 0.5 * same_worst) + + def test_mix_cuts_a_batch_only_when_the_next_frame_does_not_fit(self): + """No batch is closed early, which is what keeps batches full. + + Sorting a sliding window instead of the whole group would break this: + each window boundary closes a batch regardless of how full it is, and + an under-filled batch over-weights every frame it holds. + """ + budget = 36 + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36") + # Chop the group directly, which keeps the batches in the atom-count + # order the greedy produced them in. + batches = _chop_mixed_nloc(reader, list(range(len(reader)))) + for batch, following in pairwise(batches): + next_nloc = min(reader.frame_nlocs[index] for index in following) + self.assertGreater((len(batch) + 1) * next_nloc, budget) + + def _ragged_reader(self, batch_size="mix:36"): + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=batch_size) + reader.add_data_requirement( + [ + DataRequirementItem("force", 3, atomic=True, must=False), + DataRequirementItem("energy", 1, atomic=False, must=False), + ] + ) + reader.use_ragged_batches(True) + return reader + + def test_ragged_layout_concatenates_the_frames(self): + """The ragged layout carries the same rows on a flat, unpadded axis.""" + rectangular = LmdbDataReader( + self._lmdb_path, self._type_map, batch_size="mix:36" + ) + rectangular.add_data_requirement( + [ + DataRequirementItem("force", 3, atomic=True, must=False), + DataRequirementItem("energy", 1, atomic=False, must=False), + ] + ) + ragged = self._ragged_reader() + + # Frames 0 and 8 hold 6 and 12 atoms respectively. + indices = [0, 8] + padded = rectangular.decode_batch(indices) + flat = ragged.decode_batch(indices) + + self.assertEqual(padded["coord"].shape, (2, 12, 3)) + self.assertEqual(flat["coord"].shape, (18, 3)) + self.assertEqual(flat["atype"].shape, (18,)) + self.assertEqual(flat["force"].shape, (18, 3)) + np.testing.assert_array_equal(flat["n_node"], [6, 12]) + self.assertTrue((flat["atype"] >= 0).all(), "a ragged batch pads nothing") + # Frame-level fields keep their frame axis under either layout. + self.assertEqual(flat["energy"].shape, padded["energy"].shape) + self.assertEqual(flat["box"].shape, padded["box"].shape) + + offset = 0 + for row, count in enumerate(flat["n_node"].tolist()): + for field in ("coord", "atype", "force"): + np.testing.assert_array_equal( + padded[field][row, :count], + flat[field][offset : offset + count], + err_msg=field, + ) + offset += count + + def test_ragged_chunked_decode_matches_in_process_decode(self): + """Chunks concatenate on the flat axis, so a split decode is the same.""" + reader = self._ragged_reader() + indices = [0, 8] + layout = reader.batch_layout(indices) + chunks = [ + lmdb_data_module.decode_lmdb_batch( + reader._transaction(), + [key], + reader.frame_format, + reader._decode_config, + layout.chunk(position, position + 1), + ) + for position, key in enumerate(reader.original_keys(indices)) + ] + merged = lmdb_data_module._merge_lmdb_chunks(chunks) + expected = reader.decode_batch(indices) + self.assertEqual(sorted(merged), sorted(expected)) + for field, value in expected.items(): + np.testing.assert_array_equal(merged[field], value, err_msg=field) + + def test_ragged_packing_fills_on_real_atoms_and_keeps_the_caller_order(self): + """Without padding the budget counts real atoms, and sorting is dropped. + + Sorting exists only to keep a rectangular batch's padded width close to + its frames. A ragged batch has no width, so sorting would merely make + each optimizer step homogeneous in system size. + """ + budget = 36 + ragged = self._ragged_reader(batch_size=f"mix:{budget}") + nlocs = np.asarray(ragged.frame_nlocs) + # The fixture stores its frames in ascending atom count, so hand them + # over reversed: only then does sorting leave a visible trace. + group = list(reversed(range(len(ragged)))) + batches = _chop_mixed_nloc(ragged, group) + + for batch in batches: + if len(batch) > 1: + self.assertLessEqual(int(nlocs[batch].sum()), budget) + self.assertEqual(sorted(index for b in batches for index in b), sorted(group)) + # The frames arrive in the order they were handed over, not sorted. + self.assertEqual([index for batch in batches for index in batch], group) + + rectangular = LmdbDataReader( + self._lmdb_path, self._type_map, batch_size=f"mix:{budget}" + ) + padded = _chop_mixed_nloc(rectangular, group) + self.assertEqual( + [int(nlocs[index]) for batch in padded for index in batch], + sorted(nlocs[group].tolist()), + "the rectangular layout must still sort by atom count", + ) + + def test_mix_uses_the_fewest_batches_possible(self): + """The packing attains the minimum batch count, not merely a good one. + + The minimum is computed by dynamic programming over the sorted atom + counts. Its correctness rests on being free to take every batch + contiguous in that order: the batch holding the widest frame may be + given the widest frames outright, since its capacity ``budget // nloc`` + does not depend on which frames fill it. The recurrence then closes + each batch at its widest frame, + + steps[i] = 1 + steps[max(0, i - budget // nloc[i - 1])], + + whereas the implementation opens each batch at its narrowest. The two + must agree. + """ + for budget in (18, 36, 72): + with self.subTest(budget=budget): + reader = LmdbDataReader( + self._lmdb_path, self._type_map, batch_size=f"mix:{budget}" + ) + nlocs = sorted(reader.frame_nlocs) + steps = [0] * (len(nlocs) + 1) + for i in range(1, len(nlocs) + 1): + capacity = max(1, budget // nlocs[i - 1]) + steps[i] = 1 + steps[max(0, i - capacity)] + batches = _chop_mixed_nloc(reader, list(range(len(reader)))) + self.assertEqual(len(batches), steps[len(nlocs)]) + + def test_mix_distributed_partition_is_disjoint_and_complete(self): + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:24") + ranks = [ + DistributedLmdbBatchSampler( + reader, rank=rank, world_size=2, shuffle=True, seed=4 + ) + for rank in (0, 1) + ] + batches = [list(sampler) for sampler in ranks] + for sampler, own in zip(ranks, batches, strict=True): + self.assertEqual(len(sampler), len(own)) + indices = [i for own in batches for batch in own for i in batch] + self.assertEqual(sorted(indices), list(range(len(reader)))) + # --- LmdbTestData mixed-nloc tests --- def test_test_data_nloc_groups(self): @@ -941,7 +1389,7 @@ def test_positive_out_of_range_type_still_raises(self): class TestAutoProb(unittest.TestCase): """Test auto_prob support: frame_system_ids, compute_block_targets, - _expand_indices_by_blocks, and SameNlocBatchSampler with block_targets. + _expand_indices_by_blocks, and LmdbBatchSampler with block_targets. """ @classmethod @@ -1079,9 +1527,7 @@ def test_sampler_with_block_targets(self): nsystems=3, system_nframes=[100, 200, 300], ) - sampler = SameNlocBatchSampler( - reader, shuffle=True, block_targets=block_targets - ) + sampler = LmdbBatchSampler(reader, shuffle=True, block_targets=block_targets) all_indices = [i for batch in sampler for i in batch] self.assertGreater(len(all_indices), 600) self.assertEqual(len(set(all_indices)), 600) @@ -1092,9 +1538,12 @@ def test_sampler_allocates_block_target_across_find_signatures(self): class TwoSignatureReader: """Minimal reader exposing two one-frame availability groups.""" + mixed_nloc = False + def __init__(self): self.nloc_groups = {6: [0, 1]} self.frame_system_ids = [0, 0] + self.frame_nlocs = [6, 6] @staticmethod def group_indices_by_find_signature(indices): @@ -1106,7 +1555,7 @@ def get_batch_size_for_nloc(nloc): self.assertEqual(nloc, 6) return 1 - sampler = SameNlocBatchSampler( + sampler = LmdbBatchSampler( TwoSignatureReader(), shuffle=False, block_targets=[([0], 3)], @@ -1120,7 +1569,7 @@ def get_batch_size_for_nloc(nloc): def test_sampler_without_block_targets(self): reader = LmdbDataReader(self._lmdb_path, ["O", "H"]) - sampler = SameNlocBatchSampler(reader, shuffle=False) + sampler = LmdbBatchSampler(reader, shuffle=False) all_indices = [i for batch in sampler for i in batch] self.assertEqual(sorted(all_indices), list(range(600))) @@ -1317,11 +1766,11 @@ def test_filter_dataset_index_is_contiguous_and_live(self): reader[-1] def test_sampler_with_filter(self): - """SameNlocBatchSampler only emits retained, same-nloc frames.""" + """LmdbBatchSampler only emits retained, same-nloc frames.""" reader = LmdbDataReader( self._mixed_path, self._type_map, batch_size="filter:10" ) - sampler = SameNlocBatchSampler(reader, shuffle=False, seed=0) + sampler = LmdbBatchSampler(reader, shuffle=False, seed=0) all_batches = list(sampler) all_indices = [idx for batch in all_batches for idx in batch] @@ -1354,22 +1803,6 @@ def test_invalid_batch_size_strings_rejected(self): LmdbDataReader(self._uniform_path, self._type_map, batch_size=spec) self.assertIn("positive", str(ctx.exception)) - def test_filter_with_mixed_batch_rejected(self): - """``filter:N`` + ``mixed_batch=True`` must fail loudly. - - The mixed-batch fast path skips the per-frame nloc scan, so - filter:N cannot honour its documented ``nloc > N`` drop. - """ - with self.assertRaises(ValueError) as ctx: - LmdbDataReader( - self._mixed_path, - self._type_map, - batch_size="filter:10", - mixed_batch=True, - ) - self.assertIn("filter", str(ctx.exception)) - self.assertIn("mixed_batch", str(ctx.exception)) - def test_auto_prob_with_filter_still_works(self): """compute_block_targets + sampler survive a fully-dropped block.""" path = f"{self._tmpdir.name}/auto_prob_filter.lmdb" @@ -1399,7 +1832,7 @@ def test_auto_prob_with_filter_still_works(self): # Block 0 under-represented relative to weight → expansion needed. self.assertGreater(len(block_targets), 0) - sampler = SameNlocBatchSampler( + sampler = LmdbBatchSampler( reader, shuffle=False, seed=0, block_targets=block_targets ) all_batches = list(sampler) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 337c325fdc..1d51a6a3e1 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -1,20 +1,22 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Reusable grad-accumulation invariant harness for dpmodel loss tests. +"""Padding-mask behaviour of the dpmodel losses. -This module provides ``assert_grad_accum_invariant`` for Tasks 2-5 that -verify the loss on a padded multi-frame batch equals mean(per_frame_loss). +A batch may hold frames of unequal atom count, padded to a common width with +slots whose ``atype`` is negative. Every loss term must then reduce over the +real atoms alone. The dpmodel losses accept numpy arrays through the +array_api_compat backend. -The dpmodel losses accept numpy arrays (via the array_api_compat backend). +Harness +------- +``assert_grad_accum_invariant`` checks that the loss of a padded two-frame +batch equals the weighted mean of the frames' individual losses. The weights +express how much each frame counts: frame-level terms (energy, virial) weigh +frames equally, while per-atom terms (force, atomic energy, dos, tensor) weigh +them by their label count, so a caller passes the frames' atom counts there. -Scope / follow-ups (mixed_type padding fix, PR #5738) ----------------------------------------------------- -- The TF backend loss is not covered here and still has the mixed_type - dilution behavior; tracked in deepmodeling/deepmd-kit#5760. -- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; - tracked in deepmodeling/deepmd-kit#5761. -- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean - reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real - MAE; the grad-accum invariant is asserted by the test below. +Not covered: the TF backend, which still dilutes mixed_type losses +(deepmodeling/deepmd-kit#5760), and the pt-only ``population`` and ``denoise`` +losses (deepmodeling/deepmd-kit#5761). Constants --------- @@ -24,6 +26,7 @@ """ import numpy as np +import pytest from deepmd.dpmodel.loss.dos import ( DOSLoss, @@ -58,14 +61,24 @@ def assert_grad_accum_invariant( make_batch_A, make_batch_B, make_padded_batch, + frame_weights: tuple[float, float] = (1.0, 1.0), rtol: float = 1e-5, atol: float = 1e-6, ) -> None: - """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + """Assert that a padded batch scores the weighted mean of its frames. - The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms - padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as - processing each frame separately and averaging. + A padded batch of [frame_A (NA real atoms padded to NP) + frame_B (NB==NP + real atoms)] must yield the same loss as processing each frame separately + and combining them under the weights the term assigns to a frame: + + - **Frame-level terms** (energy, virial, global dos/cdf, global tensor, + property) carry a fixed number of labels per frame, so the frames weigh + equally and the reference is the plain mean. Pass the default weights. + - **Per-atom terms** (force, atomic energy, atomic prefactor force, atomic + dos/cdf, local tensor) carry labels in proportion to a frame's atom + count, and the loss pools them, so the reference is the atom-weighted + mean. Pass ``frame_weights=(NA, NB)``; the per-atom component count + cancels out of the ratio. Parameters ---------- @@ -79,6 +92,8 @@ def assert_grad_accum_invariant( make_padded_batch : callable Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + frame_weights : tuple[float, float] + Weight of frame A and frame B in the reference combination. rtol : float Relative tolerance for ``np.isclose``. atol : float @@ -90,7 +105,8 @@ def assert_grad_accum_invariant( loss_A = float(loss_fn(pred_A, label_A, natoms_A)) loss_B = float(loss_fn(pred_B, label_B, natoms_B)) - ref = 0.5 * (loss_A + loss_B) + weight_A, weight_B = frame_weights + ref = (weight_A * loss_A + weight_B * loss_B) / (weight_A + weight_B) loss_pad = float(loss_fn(pred_pad, label_pad, natoms_pad)) @@ -209,7 +225,9 @@ def make_B(): def make_padded(): return self._make_padded_batch(pred_A, label_A, pred_B, label_B) - assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + assert_grad_accum_invariant( + self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB) + ) def test_acdf_grad_accum_invariant(self): """Atomic cdf per-frame masked mean meets the grad-accum invariant.""" @@ -366,7 +384,9 @@ def make_padded(): NP, ) - assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + assert_grad_accum_invariant( + self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB) + ) def test_no_op_for_non_mixed(self): """All-ones mask gives same loss as no mask (non-mixed batch).""" @@ -514,6 +534,99 @@ def _padded_atom_flat(arr_A, arr_B, ncomp): ) # [2, NP] +class TestPerFrameCountSource: + """Where the per-frame atom count comes from must not change the loss. + + The extensive terms divide each frame's residual by that frame's atom + count. A padded batch states it through its mask, which the per-atom terms + also need to skip the padded rows; a concatenated batch has no padded row + and states the counts outright in ``n_node``. The two are the same number, + so the loss they produce is the same. + """ + + @staticmethod + def _loss(**kwargs): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + **kwargs, + ) + + @pytest.mark.parametrize("loss_func", ["mse", "mae"]) + @pytest.mark.parametrize("intensive", [False, True]) + @pytest.mark.parametrize("use_huber", [False, True]) + def test_mask_and_n_node_agree(self, loss_func, intensive, use_huber): + if use_huber and loss_func != "mse": + pytest.skip("huber replaces the mse branch only") + rng = np.random.default_rng(3) + nf, nloc = 3, 5 + pred, label = _full_ener_dicts( + nf, + nloc, + rng.normal(size=(nf, 1)), + rng.normal(size=(nf, 1)), + mask=np.ones((nf, nloc), dtype=np.float64), + virial=rng.normal(size=(nf, 9)), + find_virial=1.0, + ) + loss_obj = self._loss( + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + from_mask, more_mask = loss_obj.call(1.0, nloc, pred, label) + + # The same batch, described by its counts instead of by a mask. + by_count = {k: v for k, v in pred.items() if k != "mask"} + by_count["n_node"] = np.full(nf, nloc, dtype=np.int64) + from_counts, more_counts = loss_obj.call(1.0, nloc, by_count, label) + + np.testing.assert_allclose(float(from_counts), float(from_mask), rtol=0, atol=0) + assert sorted(more_counts) == sorted(more_mask) + for key, value in more_mask.items(): + np.testing.assert_allclose( + np.asarray(more_counts[key], dtype=np.float64), + np.asarray(value, dtype=np.float64), + rtol=0, + atol=0, + err_msg=key, + ) + + def test_generalized_force_refuses_a_concatenated_batch(self): + """``drdq`` is stored against a common atom axis, which is gone.""" + rng = np.random.default_rng(5) + nf, nloc = 2, 4 + pred, label = _full_ener_dicts( + nf, nloc, rng.normal(size=(nf, 1)), rng.normal(size=(nf, 1)) + ) + pred["n_node"] = np.full(nf, nloc, dtype=np.int64) + label["drdq"] = np.zeros((nf, nloc * 3, 2), dtype=np.float64) + label["find_drdq"] = 1.0 + loss_obj = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=2, + ) + with pytest.raises(NotImplementedError, match="same number of atoms"): + loss_obj.call(1.0, nloc, pred, label) + + class TestDPModelEnergyLossEnerGradAccum: """Idiom 2 (extensive) for the energy (has_e) term in EnergyLoss. @@ -703,6 +816,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -731,6 +845,28 @@ def test_huber_grad_accum(self): self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat ) + @pytest.mark.parametrize( + ("loss_func", "use_huber"), [("mae", False), ("mse", True)] + ) + def test_f_use_norm_grad_accum(self, loss_func, use_huber): + """One L2 norm per atom weighs frames the same way three components do. + + ``f_use_norm`` changes how many labels an atom contributes, so it also + changes the divisor of the pooled reduction; the frame weights it + produces must still follow the atom counts. + """ + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant( + self._make_loss(loss_func, use_huber=use_huber, f_use_norm=True), + f_A, + f_A_hat, + f_B, + f_B_hat, + ) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force loss as no mask.""" f = _rnd(NP, 3) @@ -1230,6 +1366,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -1385,6 +1522,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -1834,6 +1972,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -2165,6 +2304,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): diff --git a/source/tests/common/dpmodel/test_loss_reduction.py b/source/tests/common/dpmodel/test_loss_reduction.py index fecb0156c6..ddb59993c3 100644 --- a/source/tests/common/dpmodel/test_loss_reduction.py +++ b/source/tests/common/dpmodel/test_loss_reduction.py @@ -14,15 +14,12 @@ class TestMaskedAtomMean: - """Idiom 1: per-atom masked mean over ncomp components, averaged over frames.""" + """Idiom 1: mean of a per-atom contribution over the batch's real labels.""" def _ref(self, elem, maskf, ncomp): - # reference reduction, numpy - nf = elem.shape[0] + # reference reduction, numpy: pool over every real label of the batch masked = elem * maskf[:, :, None] - pfs = masked.reshape(nf, -1).sum(axis=-1) - pfd = maskf.sum(axis=-1) * ncomp - return (pfs / pfd).mean() + return masked.sum() / (maskf.sum() * ncomp) @pytest.mark.parametrize("ncomp", [1, 3]) # atom-energy (1) and force (3) def test_numpy_matches_reference(self, ncomp) -> None: @@ -32,7 +29,36 @@ def test_numpy_matches_reference(self, ncomp) -> None: got = masked_atom_mean(elem, maskf, ncomp) np.testing.assert_allclose(got, self._ref(elem, maskf, ncomp), rtol=0, atol=0) - def test_torch_autograd_and_bit_identical(self) -> None: + def test_uniform_atom_count_batch_is_a_frame_mean(self) -> None: + """Pooling and the per-frame mean agree when every frame has one size. + + This identity is what lets the pooled convention apply unconditionally: + a batch whose frames share an atom count keeps the reduction it had + before mixed-nloc batches existed, with no special case for it. + """ + rng = np.random.default_rng(5) + elem = rng.random((3, 4, 3)) + maskf = np.tile(np.array([1.0, 1.0, 1.0, 0.0]), (3, 1)) + per_frame = (elem * maskf[:, :, None]).reshape(3, -1).sum(axis=-1) / ( + maskf.sum(axis=-1) * 3 + ) + np.testing.assert_allclose( + masked_atom_mean(elem, maskf, 3), per_frame.mean(), rtol=1e-14, atol=0 + ) + + def test_frames_weigh_by_their_label_count(self) -> None: + """A frame's weight is its share of the batch's labels, not ``1 / nf``.""" + # Frame 0 keeps one real atom, frame 1 keeps three. + maskf = np.array([[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]]) + elem = np.zeros((2, 4, 3)) + elem[0, :1] = 2.0 + elem[1, :3] = 6.0 + # Pooled: (1*3*2 + 3*3*6) / (4*3); a frame mean would give (2+6)/2 = 4. + np.testing.assert_allclose( + masked_atom_mean(elem, maskf, 3), (6.0 + 54.0) / 12.0, rtol=0, atol=0 + ) + + def test_torch_autograd_and_matches_numpy(self) -> None: elem_np = np.random.default_rng(1).random((2, 4, 3)) maskf_np = np.array([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]]) elem = torch.tensor(elem_np, requires_grad=True) @@ -40,34 +66,24 @@ def test_torch_autograd_and_bit_identical(self) -> None: out = masked_atom_mean(elem, maskf, 3) out.backward() assert elem.grad is not None - # bit-identical to torch-native inline form - en = torch.tensor(elem_np) - mn = torch.tensor(maskf_np) - pfs = (en * mn[:, :, None]).reshape(2, -1).sum(dim=-1) - pfd = mn.sum(dim=-1) * 3 - ref = torch.mean(pfs / pfd) - assert out.item() == ref.item() - - def test_all_padding_frame_is_not_nan(self) -> None: - # a frame with zero real atoms has per_frame_dof == 0; the ratio must - # not become 0/0 = NaN (an independent finiteness invariant -- the - # reference formula shares the bug, so equality checks cannot catch it) + np.testing.assert_allclose( + out.item(), self._ref(elem_np, maskf_np, 3), rtol=1e-14, atol=0 + ) + + def test_all_padding_batch_is_not_nan(self) -> None: + # With the pooled reduction only a batch without a single real atom + # drives the denominator to zero; the ratio must not become 0/0 = NaN. elem = np.random.default_rng(3).random((2, 4, 3)) - maskf = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) + maskf = np.zeros((2, 4)) got = masked_atom_mean(elem, maskf, 3) assert np.isfinite(got) - # the all-padding frame contributes a neutral 0, so the result is the - # first frame's masked mean divided by the number of frames - pfs0 = (elem[0] * maskf[0][:, None]).reshape(-1).sum() - expected = (pfs0 / (maskf[0].sum() * 3)) / 2 - np.testing.assert_allclose(got, expected, rtol=0, atol=0) + np.testing.assert_allclose(got, 0.0, rtol=0, atol=0) - def test_all_padding_frame_torch_grad_is_not_nan(self) -> None: + def test_all_padding_batch_torch_grad_is_not_nan(self) -> None: # the guard must keep both the value and the autograd gradient finite elem_np = np.random.default_rng(4).random((2, 4, 3)) - maskf_np = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) elem = torch.tensor(elem_np, requires_grad=True) - maskf = torch.tensor(maskf_np) + maskf = torch.zeros((2, 4), dtype=torch.float64) out = masked_atom_mean(elem, maskf, 3) out.backward() assert torch.isfinite(out).item() diff --git a/source/tests/consistent/test_lmdb_data.py b/source/tests/consistent/test_lmdb_data.py index 6e9cecee52..c4f35563ee 100644 --- a/source/tests/consistent/test_lmdb_data.py +++ b/source/tests/consistent/test_lmdb_data.py @@ -14,6 +14,8 @@ from deepmd.dpmodel.utils.lmdb_data import ( LmdbDataReader, + collate_lmdb_frames, + resolve_per_atom_keys, ) try: @@ -178,7 +180,7 @@ def test_same_properties(self): self.assertEqual(self._reader.batch_sizes, self._ds.batch_sizes) self.assertEqual(self._reader.nframes, self._ds.nframes) self.assertEqual(self._reader.mixed_type, self._ds.mixed_type) - self.assertEqual(self._reader.mixed_batch, self._ds.mixed_batch) + self.assertEqual(self._reader.mixed_nloc, self._ds.mixed_nloc) def test_data_requirement(self): req = [ @@ -219,8 +221,30 @@ def test_mixed_nloc_same_properties(self): ds = LmdbDataset(path, self._type_map, batch_size=2) self.assertEqual(reader.nframes, ds.nframes) self.assertEqual(reader.batch_sizes, ds.batch_sizes) - self.assertEqual(reader.mixed_batch, ds.mixed_batch) - self.assertFalse(reader.mixed_batch) + self.assertEqual(reader.mixed_nloc, ds.mixed_nloc) + self.assertFalse(reader.mixed_nloc) + + def test_mix_batch_size_same_padded_batch(self): + """Both decode paths pad a mixed-nloc batch to the same arrays.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = _create_mixed_nloc_lmdb(f"{tmpdir}/mixed.lmdb") + reader = LmdbDataReader(path, self._type_map, batch_size="mix:24") + self.assertTrue(reader.mixed_nloc) + # _create_mixed_nloc_lmdb interleaves 6-, 9- and 12-atom frames. + indices = [0, 4, 8] + per_atom_keys = resolve_per_atom_keys( + reader[indices[0]], reader.decode_config + ) + expected = collate_lmdb_frames( + [reader[index] for index in indices], per_atom_keys + ) + actual = reader.decode_batch(indices) + self.assertEqual(tuple(actual), tuple(expected)) + for key, expected_value in expected.items(): + if isinstance(expected_value, np.ndarray): + np.testing.assert_array_equal(actual[key], expected_value) + else: + self.assertEqual(actual[key], expected_value) if __name__ == "__main__": diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 92a1cc14a6..9b1028b4cb 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -1692,6 +1692,20 @@ def _frame( ) return coord, atype, spin, box + def _mixed_padded_batch( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build two frames whose shorter member carries phantom padding.""" + coord, atype, spin, box = self._frame() + coord = coord.repeat(2, 1, 1) + atype = atype.repeat(2, 1) + spin = spin.repeat(2, 1, 1) + box = box.repeat(2, 1) + atype[1, -2:] = -1 + coord[1, -2:] = 0.0 + spin[1, -2:] = 0.0 + return coord, atype, spin, box + def test_zbl_change_out_bias_is_invariant_for_self_labels(self) -> None: """Native-spin statistics consume spin and the complete ZBL energy.""" model = self._build_model(bridging_method="ZBL") @@ -1796,6 +1810,33 @@ def test_joint_rotation_equivariance(self) -> None: rtol=1e-6, ) + def test_phantom_atoms_are_excluded_from_magnetic_mask(self) -> None: + """Mixed-nloc padding never enters the per-type spin lookup.""" + model = self._build_model() + coord, atype, spin, box = self._mixed_padded_batch() + + out = model(coord, atype, spin, box=box) + expected_mask = ((atype == 0) & (atype >= 0)).unsqueeze(-1) + torch.testing.assert_close(out["mask_mag"], expected_mask) + self.assertTrue(torch.all(out["force"][1, -2:] == 0.0)) + self.assertTrue(torch.all(out["force_mag"][1, -2:] == 0.0)) + + lower = model._attach_spin_masks( + { + "energy_derv_r_mag": torch.zeros( + 2, + atype.shape[1], + 1, + 3, + dtype=coord.dtype, + device=coord.device, + ) + }, + atype=atype, + nall=atype.shape[1], + ) + torch.testing.assert_close(lower["mask_mag"], expected_mask) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_export_matches_forward(self) -> None: """The traced ``.pt2`` export reduces to the public forward. @@ -1904,7 +1945,7 @@ def test_ener_spin_loss_smoke(self) -> None: @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_compile_matches_eager(self) -> None: """The compiled native-spin path matches eager force and magnetic force.""" - coord, atype, spin, box = self._frame() + coord, atype, spin, box = self._mixed_padded_batch() model_eager = self._build_model(use_compile=False) model_cmp = self._build_model(use_compile=True) model_cmp.load_state_dict(model_eager.state_dict()) @@ -1922,6 +1963,8 @@ def test_compile_matches_eager(self) -> None: rtol=1.0e-6, msg=f"native-spin compile mismatch on {key}", ) + torch.testing.assert_close(out_c["mask_mag"], out_e["mask_mag"]) + self.assertTrue(torch.all(~out_c["mask_mag"][1, -2:])) @staticmethod def _extended_spin_inputs( @@ -2238,6 +2281,147 @@ def test_zbl_respects_exclusions(self) -> None: ) +class TestSeZMPhantomAtoms(unittest.TestCase): + """SeZM must ignore the phantom atoms that pad a mixed-nloc batch. + + A mixed-nloc LMDB batch is rectangular: frames shorter than the batch-wide + atom count are padded with slots carrying ``atype = -1`` and zero + coordinates. Those slots stand for no physical site, so the padded batch + has to reproduce, frame by frame, what the unpadded frames give on their + own. + """ + + def setUp(self) -> None: + self.device = env.DEVICE + torch.manual_seed(2024) + + def _build_model(self) -> SeZMModel: + params = { + "type": "SeZM", + "type_map": ["O", "H"], + "descriptor": { + "type": "SeZM", + "sel": [20, 20], + "rcut": 4.0, + "channels": 4, + "n_focus": 1, + "focus_compete": False, + "n_radial": 3, + "radial_mlp": [6], + "use_env_seed": False, + "l_schedule": [1, 0], + "mmax": 1, + "so2_norm": False, + "so2_layers": 1, + "n_atten_head": 0, + "sandwich_norm": [True, False, True, False], + "ffn_neurons": 8, + "ffn_blocks": 1, + "mlp_bias": True, + "layer_scale": True, + "use_amp": False, + "activation_function": "silu", + "glu_activation": True, + "precision": "float64", + "seed": 7, + }, + "fitting_net": { + "neuron": [8], + "activation_function": "silu", + "precision": "float64", + "seed": 7, + }, + "use_compile": False, + } + model = get_sezm_model(params).to(self.device) + torch.manual_seed(1234) + with torch.no_grad(): + for param in model.parameters(): + param.copy_(torch.randn_like(param) * 0.1) + return model.eval() + + @staticmethod + def _make_frames() -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: + """Three frames of unequal atom count sharing one cubic cell.""" + rng = np.random.default_rng(0) + nlocs = (4, 7, 3) + coords = [rng.uniform(0.0, 6.0, (nloc, 3)) for nloc in nlocs] + atypes = [rng.integers(0, 2, nloc) for nloc in nlocs] + return coords, atypes, (np.eye(3) * 10.0).reshape(9) + + def _pad( + self, coords: list[np.ndarray], atypes: list[np.ndarray] + ) -> tuple[torch.Tensor, torch.Tensor]: + pad_nloc = max(len(a) for a in atypes) + coord = np.zeros((len(atypes), pad_nloc, 3)) + atype = np.full((len(atypes), pad_nloc), -1, dtype=np.int64) + for index, (frame_coord, frame_atype) in enumerate( + zip(coords, atypes, strict=True) + ): + coord[index, : len(frame_atype)] = frame_coord + atype[index, : len(frame_atype)] = frame_atype + return ( + torch.tensor(coord, dtype=torch.float64, device=self.device), + torch.tensor(atype, device=self.device), + ) + + def test_padded_batch_matches_unpadded_frames(self) -> None: + """Energy and real-atom forces are unchanged by the padding.""" + model = self._build_model() + coords, atypes, box = self._make_frames() + coord, atype = self._pad(coords, atypes) + batched = model( + coord, + atype, + box=torch.tensor( + np.tile(box, (len(atypes), 1)), + dtype=torch.float64, + device=self.device, + ), + ) + + for index, (frame_coord, frame_atype) in enumerate( + zip(coords, atypes, strict=True) + ): + alone = model( + torch.tensor( + frame_coord[None], dtype=torch.float64, device=self.device + ), + torch.tensor(frame_atype[None], device=self.device), + box=torch.tensor(box[None], dtype=torch.float64, device=self.device), + ) + torch.testing.assert_close( + batched["energy"].reshape(-1)[index], + alone["energy"].reshape(-1)[0], + atol=1.0e-12, + rtol=1.0e-12, + ) + torch.testing.assert_close( + batched["force"][index, : len(frame_atype)], + alone["force"][0], + atol=1.0e-12, + rtol=1.0e-12, + ) + + def test_phantom_atoms_carry_no_force(self) -> None: + """Padded slots stay at exactly zero force, so no gradient reaches them.""" + model = self._build_model() + coords, atypes, box = self._make_frames() + coord, atype = self._pad(coords, atypes) + force = model( + coord, + atype, + box=torch.tensor( + np.tile(box, (len(atypes), 1)), + dtype=torch.float64, + device=self.device, + ), + )["force"] + phantom = atype < 0 + self.assertTrue(bool(phantom.any()), "fixture must exercise padding") + self.assertTrue(bool(torch.all(force[phantom] == 0.0))) + + class TestSeZMModelModes(unittest.TestCase): """Targeted regression tests for SeZM `ener` / `dens` mode routing.""" diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index 4b253aa1ae..2010b57016 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Unit tests for LmdbDataset (PyTorch wrapper) and related PT-specific features. -Pure dpmodel tests (LmdbDataReader, LmdbTestData, SameNlocBatchSampler, type_map +Pure dpmodel tests (LmdbDataReader, LmdbTestData, LmdbBatchSampler, type_map remapping, auto_prob) live in source/tests/common/dpmodel/test_lmdb_data.py. Consistency tests (dpmodel vs pt) live in source/tests/consistent/test_lmdb_data.py. """ @@ -17,9 +17,11 @@ ) from deepmd.dpmodel.utils.lmdb_data import ( _ENV_CACHE, - DistributedSameNlocBatchSampler, + PHANTOM_ATOM_TYPE, + DistributedLmdbBatchSampler, + LmdbBatchSampler, LmdbDataReader, - SameNlocBatchSampler, + LmdbDecodeConfig, _decode_frame, _read_metadata, _remap_keys, @@ -350,9 +352,7 @@ def test_batch_iteration(self, lmdb_dir): ) with torch.device("cpu"): - dl = DataLoader( - ds, batch_size=2, shuffle=False, collate_fn=_collate_lmdb_batch - ) + dl = DataLoader(ds, batch_size=2, shuffle=False, collate_fn=ds._collate) batch = next(iter(dl)) assert batch["coord"].shape == (2, 6, 3) assert batch["energy"].shape == (2, 1) @@ -368,7 +368,7 @@ def test_inner_dataloader(self, lmdb_dir): def test_parallel_batch_loader_has_finite_epoch(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) - sampler = SameNlocBatchSampler(ds._reader, shuffle=False) + sampler = LmdbBatchSampler(ds._reader, shuffle=False) loader = LmdbBatchDataLoader( ds, sampler, @@ -389,13 +389,13 @@ def test_parallel_loaders_share_pool_for_same_dataset(self, lmdb_dir): ) first = LmdbBatchDataLoader( first_data, - SameNlocBatchSampler(first_data._reader, shuffle=True, seed=1), + LmdbBatchSampler(first_data._reader, shuffle=True, seed=1), pin_memory=False, num_workers=2, ) second = LmdbBatchDataLoader( second_data, - SameNlocBatchSampler(second_data._reader, shuffle=True, seed=2), + LmdbBatchSampler(second_data._reader, shuffle=True, seed=2), pin_memory=False, num_workers=2, ) @@ -415,7 +415,7 @@ def test_small_batch_stays_synchronous(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) loader = LmdbBatchDataLoader( ds, - SameNlocBatchSampler(ds._reader, shuffle=False), + LmdbBatchSampler(ds._reader, shuffle=False), pin_memory=False, num_workers=4, ) @@ -430,7 +430,7 @@ def test_partial_successor_is_deferred(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=4) loader = LmdbBatchDataLoader( ds, - SameNlocBatchSampler(ds._reader, shuffle=False), + LmdbBatchSampler(ds._reader, shuffle=False), pin_memory=False, num_workers=4, ) @@ -450,7 +450,7 @@ def test_requirements_freeze_after_batch_read(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) loader = LmdbBatchDataLoader( ds, - SameNlocBatchSampler(ds._reader, shuffle=False), + LmdbBatchSampler(ds._reader, shuffle=False), pin_memory=False, num_workers=0, ) @@ -468,12 +468,73 @@ def test_full_epoch(self, lmdb_dir): ) with torch.device("cpu"): - dl = DataLoader( - ds, batch_size=3, shuffle=False, collate_fn=_collate_lmdb_batch - ) + dl = DataLoader(ds, batch_size=3, shuffle=False, collate_fn=ds._collate) total_frames = sum(batch["coord"].shape[0] for batch in dl) assert total_frames == 10 + def test_loss_ignores_phantom_atoms(self, lmdb_dir, monkeypatch): + """Padding a batch with phantom atoms leaves its loss untouched. + + This closes the loop from the loader to the loss: the slots a + mixed-nloc batch adds must enter neither the energy term, which + normalizes by atom count, nor the per-atom force mean. + """ + monkeypatch.setattr("deepmd.pt.loss.ener.env.DEVICE", torch.device("cpu")) + ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=3) + ds.add_data_requirement( + [ + DataRequirementItem("energy", 1, atomic=False, must=False), + DataRequirementItem("force", 3, atomic=True, must=False), + ] + ) + loss_module = EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=1.0, + limit_pref_f=1.0, + ) + + def score(batch): + """Loss of a constant-zero prediction against this batch's labels.""" + + def zero_model(**kwargs): + return { + "energy": torch.zeros_like(batch["energy"]), + "force": torch.zeros_like(batch["force"]), + } + + _, loss, _ = loss_module( + {"atype": batch["atype"]}, + zero_model, + { + "energy": batch["energy"], + "find_energy": batch["find_energy"], + "force": batch["force"], + "find_force": batch["find_force"], + }, + natoms=int(batch["atype"].shape[-1]), + learning_rate=1.0, + ) + return float(loss) + + with torch.device("cpu"): + batch = ds._collate([ds[index] for index in range(3)]) + nframes, nloc = batch["atype"].shape + padded = dict(batch) + padded["atype"] = torch.cat( + [ + batch["atype"], + torch.full((nframes, 2), PHANTOM_ATOM_TYPE, dtype=torch.int64), + ], + dim=1, + ) + padded["force"] = torch.cat( + [batch["force"], torch.zeros((nframes, 2, 3), dtype=torch.float64)], + dim=1, + ) + assert score(padded) == pytest.approx(score(batch), rel=1e-12) + def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatch): """Default-filled labels must never share a scalar flag with real ones.""" monkeypatch.setattr("deepmd.pt.loss.ener.env.DEVICE", torch.device("cpu")) @@ -487,7 +548,7 @@ def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatc ] ) - sampler = SameNlocBatchSampler(ds._reader, shuffle=True, seed=11) + sampler = LmdbBatchSampler(ds._reader, shuffle=True, seed=11) batches = list(sampler) assert len(sampler) == len(batches) == 2 for indices in batches: @@ -496,7 +557,7 @@ def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatc distributed_batches = [] for rank in range(2): - distributed = DistributedSameNlocBatchSampler( + distributed = DistributedLmdbBatchSampler( ds._reader, rank=rank, world_size=2, @@ -594,6 +655,16 @@ def test_unrequested_labels_form_homogeneous_batches(self, tmp_path): # ============================================================ +_BARE_DECODE_CONFIG = LmdbDecodeConfig( + ntypes=2, natoms=0, type_remap=None, data_requirements={} +) + + +def _collate(frames): + """Collate hand-built frames with no registered data requirements.""" + return _collate_lmdb_batch(frames, _BARE_DECODE_CONFIG) + + class TestCollate: """Test collate function.""" @@ -613,7 +684,7 @@ def test_collate_basic(self): "fid": 1, }, ] - batch = _collate_lmdb_batch(frames) + batch = _collate(frames) assert batch["coord"].shape == (2, 4, 3) assert batch["fid"] == [0, 1] assert batch["sid"] == 0 @@ -623,14 +694,14 @@ def test_collate_skips_type(self): {"coord": np.zeros((2, 3)), "type": np.array([0, 1])}, {"coord": np.zeros((2, 3)), "type": np.array([0, 1])}, ] - assert "type" not in _collate_lmdb_batch(frames) + assert "type" not in _collate(frames) def test_collate_none_values(self): frames = [ {"coord": np.zeros((2, 3)), "box": None}, {"coord": np.zeros((2, 3)), "box": None}, ] - assert _collate_lmdb_batch(frames)["box"] is None + assert _collate(frames)["box"] is None def test_collate_rejects_mixed_find_flags(self): frames = [ @@ -638,7 +709,36 @@ def test_collate_rejects_mixed_find_flags(self): {"coord": np.zeros((2, 3)), "find_energy": 0.0}, ] with pytest.raises(ValueError, match="mixes 'find_energy' values"): - _collate_lmdb_batch(frames) + _collate(frames) + + def test_collate_pads_the_atom_axis(self): + """Frames of different atom counts stack into one padded batch.""" + frames = [ + { + "coord": np.ones((2, 3)), + "atype": np.zeros(2, dtype=np.int64), + "energy": np.array([1.0]), + }, + { + "coord": np.ones((4, 3)), + "atype": np.zeros(4, dtype=np.int64), + "energy": np.array([2.0]), + }, + ] + batch = _collate(frames) + assert batch["coord"].shape == (2, 4, 3) + assert batch["atype"].shape == (2, 4) + # The short frame keeps its two atoms and gains two phantom slots. + assert batch["atype"][0].tolist() == [ + 0, + 0, + PHANTOM_ATOM_TYPE, + PHANTOM_ATOM_TYPE, + ] + assert batch["atype"][1].tolist() == [0, 0, 0, 0] + assert torch.all(batch["coord"][0, 2:] == 0) + # Frame-level fields keep their own shape. + assert batch["energy"].shape == (2, 1) # ============================================================ @@ -743,15 +843,15 @@ def multi_nloc_lmdb(tmp_path): return lmdb_path -class TestDistributedSameNlocBatchSampler: - """Test DistributedSameNlocBatchSampler (pure logic, no torch.distributed).""" +class TestDistributedLmdbBatchSampler: + """Test DistributedLmdbBatchSampler (pure logic, no torch.distributed).""" def test_disjoint_batches(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=1) - s0 = DistributedSameNlocBatchSampler( + s0 = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) - s1 = DistributedSameNlocBatchSampler( + s1 = DistributedLmdbBatchSampler( reader, rank=1, world_size=2, shuffle=True, seed=42 ) frames0 = {i for batch in s0 for i in batch} @@ -760,10 +860,10 @@ def test_disjoint_batches(self, multi_nloc_lmdb): def test_covers_all_frames(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) - s0 = DistributedSameNlocBatchSampler( + s0 = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) - s1 = DistributedSameNlocBatchSampler( + s1 = DistributedLmdbBatchSampler( reader, rank=1, world_size=2, shuffle=True, seed=42 ) all_frames = {i for batch in s0 for i in batch} | { @@ -775,9 +875,9 @@ def test_len(self, multi_nloc_lmdb): import math reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) - total = len(SameNlocBatchSampler(reader, shuffle=False)) + total = len(LmdbBatchSampler(reader, shuffle=False)) samplers = [ - DistributedSameNlocBatchSampler( + DistributedLmdbBatchSampler( reader, rank=rank, world_size=4, shuffle=False, seed=0 ) for rank in range(4) @@ -787,17 +887,17 @@ def test_len(self, multi_nloc_lmdb): def test_deterministic(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) - s1 = DistributedSameNlocBatchSampler( + s1 = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) - s2 = DistributedSameNlocBatchSampler( + s2 = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) assert list(s1) == list(s2) def test_set_epoch_changes_order(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) - s = DistributedSameNlocBatchSampler( + s = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) s.set_epoch(0) @@ -810,12 +910,12 @@ def test_single_gpu_fallback(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) single = { i - for batch in SameNlocBatchSampler(reader, shuffle=True, seed=42) + for batch in LmdbBatchSampler(reader, shuffle=True, seed=42) for i in batch } dist = { i - for batch in DistributedSameNlocBatchSampler( + for batch in DistributedLmdbBatchSampler( reader, rank=0, world_size=1, shuffle=True, seed=42 ) for i in batch @@ -824,7 +924,7 @@ def test_single_gpu_fallback(self, multi_nloc_lmdb): def test_same_nloc_per_batch(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) - s = DistributedSameNlocBatchSampler( + s = DistributedLmdbBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) for batch in s: @@ -936,14 +1036,14 @@ def test_distributed_len_includes_auto_prob_expansion(self, auto_prob_lmdb): auto_prob_style="prob_sys_size;0:1:0.5;1:3:0.5", ) global_batches = len(ds._batch_sampler) - dist_sampler_rank0 = DistributedSameNlocBatchSampler( + dist_sampler_rank0 = DistributedLmdbBatchSampler( ds._reader, rank=0, world_size=2, shuffle=False, block_targets=ds._block_targets, ) - dist_sampler_rank1 = DistributedSameNlocBatchSampler( + dist_sampler_rank1 = DistributedLmdbBatchSampler( ds._reader, rank=1, world_size=2, @@ -955,28 +1055,26 @@ def test_distributed_len_includes_auto_prob_expansion(self, auto_prob_lmdb): assert len(dist_sampler_rank0) == len(list(dist_sampler_rank0)) assert len(dist_sampler_rank1) == len(list(dist_sampler_rank1)) - def test_distributed_len_reuses_cached_total(self, auto_prob_lmdb, monkeypatch): - calls = 0 - real_sampler = lmdb_data.SameNlocBatchSampler + def test_distributed_builds_batches_once_per_epoch( + self, auto_prob_lmdb, monkeypatch + ): + """Batch construction is shared by ``__len__`` and ``__iter__``.""" + builds = 0 + real_build = lmdb_data._build_all_batches - class CountingSameNlocBatchSampler(real_sampler): - def __init__(self, *args, **kwargs): - nonlocal calls - calls += 1 - super().__init__(*args, **kwargs) + def counting_build(*args, **kwargs): + nonlocal builds + builds += 1 + return real_build(*args, **kwargs) - monkeypatch.setattr( - lmdb_data, - "SameNlocBatchSampler", - CountingSameNlocBatchSampler, - ) + monkeypatch.setattr(lmdb_data, "_build_all_batches", counting_build) ds = LmdbDataset( auto_prob_lmdb, type_map=["O", "H"], batch_size=4, auto_prob_style="prob_sys_size;0:1:0.5;1:3:0.5", ) - dist_sampler = DistributedSameNlocBatchSampler( + dist_sampler = DistributedLmdbBatchSampler( ds._reader, rank=1, world_size=2, @@ -984,10 +1082,17 @@ def __init__(self, *args, **kwargs): block_targets=ds._block_targets, ) - assert calls == 1 + assert builds == 0 expected_len = (len(ds._batch_sampler) + 1) // 2 + builds = 0 assert len(dist_sampler) == expected_len - assert calls == 1 + assert len(list(dist_sampler)) == expected_len + assert builds == 1 + + # A new epoch reshuffles, and so must rebuild exactly once more. + dist_sampler.set_epoch(1) + assert len(list(dist_sampler)) == len(dist_sampler) + assert builds == 2 class TestMergeLmdbSystemIds: @@ -1151,6 +1256,70 @@ def fail() -> None: assert loader.closed +def test_numb_epoch_counts_mixed_nloc_batches(multi_nloc_lmdb, tmp_path, monkeypatch): + """One epoch of a ``mix:N`` dataset is one pass over its padded batches. + + The batch count of a mixed-nloc pass depends on how the shuffle groups + atom counts, so an epoch can only be measured on the sampler the trainer + will actually iterate -- not on a nominal batch size. + """ + from deepmd.pt.entrypoints.main import ( + get_trainer, + ) + from deepmd.utils.argcheck import ( + normalize, + ) + from deepmd.utils.compat import ( + update_deepmd_input, + ) + + config = { + "model": { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 8], + "axis_neuron": 4, + "precision": "float64", + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 1}, + "data_stat_nbatch": 1, + }, + "learning_rate": { + "type": "exp", + "decay_steps": 50, + "start_lr": 1e-3, + "stop_lr": 1e-8, + }, + "loss": {"type": "ener", "start_pref_e": 1.0, "start_pref_f": 1.0}, + "training": { + "training_data": { + "systems": multi_nloc_lmdb, + "batch_size": "mix:24", + }, + "numb_epoch": 3, + "seed": 10, + "disp_file": str(tmp_path / "lcurve.out"), + "disp_freq": 100, + "save_freq": 100, + }, + } + monkeypatch.chdir(tmp_path) + config = normalize(update_deepmd_input(config, warning=False)) + trainer = get_trainer(config) + try: + batches_per_epoch = len(trainer.training_dataloader) + assert trainer.training_dataloader.dataset.mixed_nloc + assert batches_per_epoch > 0 + assert trainer.num_steps == 3 * batches_per_epoch + finally: + trainer.training_dataloader.close() + + @pytest.fixture def multitask_lmdb_setup(tmp_path): """Create two LMDB datasets and a multitask training config.""" diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 094bcfeadb..927c9ba0cd 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -1,28 +1,26 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Tests for mixed_type loss padding-mask support in the pt backend. +"""Padding-mask behaviour of the pt backend losses. -Task 1: verify that TaskLoss._inject_atom_mask correctly recovers the per-atom -mask from atype (ghost atoms have atype < 0) so that later tasks can exclude -them from loss reductions. +A batch may hold frames of unequal atom count, padded to a common width with +slots whose ``atype`` is negative. ``TaskLoss._inject_atom_mask`` recovers the +per-atom mask from ``atype``, and every loss term must then reduce over the +real atoms alone. Harness ------- -assert_grad_accum_invariant -- reusable by Tasks 2-5 to check the - grad-accumulation invariant: loss on a padded multi-frame batch must equal - mean_over_frames(per_frame_loss). - -Scope / follow-ups (mixed_type padding fix, PR #5738) ----------------------------------------------------- -- The TF backend loss is not covered here and still has the mixed_type - dilution behavior; tracked in deepmodeling/deepmd-kit#5760. -- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; - tracked in deepmodeling/deepmd-kit#5761. -- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean - reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real - MAE; the grad-accum invariant is asserted by the test below. +``assert_grad_accum_invariant`` checks that the loss of a padded two-frame +batch equals the weighted mean of the frames' individual losses. The weights +express how much each frame counts: frame-level terms (energy, virial) weigh +frames equally, while per-atom terms (force, atomic energy, dos, tensor) weigh +them by their label count, so a caller passes the frames' atom counts there. + +Not covered: the TF backend, which still dilutes mixed_type losses +(deepmodeling/deepmd-kit#5760), and the pt-only ``population`` and ``denoise`` +losses (deepmodeling/deepmd-kit#5761). """ import numpy as np +import pytest import torch from deepmd.pt.loss.dos import ( @@ -64,14 +62,24 @@ def assert_grad_accum_invariant( make_batch_A, make_batch_B, make_padded_batch, + frame_weights: tuple[float, float] = (1.0, 1.0), rtol: float = 1e-5, atol: float = 1e-6, ) -> None: - """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + """Assert that a padded batch scores the weighted mean of its frames. - The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms - padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as - processing each frame separately and averaging. + A padded batch of [frame_A (NA real atoms padded to NP) + frame_B (NB==NP + real atoms)] must yield the same loss as processing each frame separately + and combining them under the weights the term assigns to a frame: + + - **Frame-level terms** (energy, virial, global dos/cdf, global tensor, + property) carry a fixed number of labels per frame, so the frames weigh + equally and the reference is the plain mean. Pass the default weights. + - **Per-atom terms** (force, atomic energy, atomic prefactor force, atomic + dos/cdf, local tensor) carry labels in proportion to a frame's atom + count, and the loss pools them, so the reference is the atom-weighted + mean. Pass ``frame_weights=(NA, NB)``; the per-atom component count + cancels out of the ratio. Parameters ---------- @@ -84,6 +92,8 @@ def assert_grad_accum_invariant( make_padded_batch : callable Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + frame_weights : tuple[float, float] + Weight of frame A and frame B in the reference combination. rtol : float Relative tolerance for ``torch.isclose``. atol : float @@ -95,7 +105,8 @@ def assert_grad_accum_invariant( loss_A = loss_fn(pred_A, label_A, natoms_A) loss_B = loss_fn(pred_B, label_B, natoms_B) - ref = 0.5 * (loss_A + loss_B) + weight_A, weight_B = frame_weights + ref = (weight_A * loss_A + weight_B * loss_B) / (weight_A + weight_B) loss_pad = loss_fn(pred_pad, label_pad, natoms_pad) @@ -216,7 +227,9 @@ def make_padded(): NP, ) - assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + assert_grad_accum_invariant( + self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB) + ) def test_no_op_for_non_mixed(self): """All-ones mask gives same loss as no mask (non-mixed batch).""" @@ -404,7 +417,9 @@ def make_padded(): NP, ) - assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + assert_grad_accum_invariant( + self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB) + ) def test_no_op_for_non_mixed(self): """All-ones mask gives same loss as no mask (non-mixed batch).""" @@ -826,6 +841,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -854,6 +870,28 @@ def test_huber_grad_accum(self): self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat ) + @pytest.mark.parametrize( + ("loss_func", "use_huber"), [("mae", False), ("mse", True)] + ) + def test_f_use_norm_grad_accum(self, loss_func, use_huber): + """One L2 norm per atom weighs frames the same way three components do. + + ``f_use_norm`` changes how many labels an atom contributes, so it also + changes the divisor of the pooled reduction; the frame weights it + produces must still follow the atom counts. + """ + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant( + self._make_loss(loss_func, use_huber=use_huber, f_use_norm=True), + f_A, + f_A_hat, + f_B, + f_B_hat, + ) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force loss as no mask.""" f = _t(NP, 3) @@ -1036,6 +1074,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -1159,6 +1198,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -1592,6 +1632,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): @@ -1867,6 +1908,7 @@ def make_padded(): make_A, make_B, make_padded, + frame_weights=(NA, NB), ) def test_mse_grad_accum(self): diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 641804b1cd..fcc93e547b 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -529,6 +529,32 @@ def test_numb_epoch_counts_passes_over_the_lmdb(self) -> None: # train.lmdb holds eight frames, read one frame per batch. self.assertEqual(trainer.num_steps, 2 * 8) + def test_numb_epoch_counts_mixed_nloc_batches(self) -> None: + """One epoch of a ``mix:N`` dataset is one pass over its padded batches. + + The batch count of a mixed-nloc pass depends on how the shuffle groups + atom counts, so an epoch can only be measured on the sampler the + trainer will actually draw from, not on a nominal batch size. + """ + config = self._make_lmdb_config() + del config["training"]["numb_steps"] + config["training"]["numb_epoch"] = 3.0 + config["training"]["training_data"]["systems"] = self.mixed_lmdb_path + config["training"]["training_data"]["batch_size"] = "mix:27" + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + trainer = get_trainer(config) + finally: + os.chdir(cwd) + + data = trainer.training_data + self.assertTrue(data._reader.mixed_nloc) + self.assertEqual(trainer.num_steps, 3 * data.nbatches[0]) + def test_training_closes_parallel_lmdb_pipeline(self) -> None: """Trainer shutdown releases spawned decoder processes.""" config = self._make_lmdb_config(numb_steps=2) @@ -567,6 +593,168 @@ def test_mixed_nloc_statistics_and_training(self) -> None: finally: os.chdir(cwd) + def test_mix_batch_size_trains_on_padded_batches(self) -> None: + """``mix:N`` reaches the trainer and its padded batches train.""" + config = self._make_lmdb_config(numb_steps=2) + config["model"]["data_stat_nbatch"] = 10 + config["training"]["training_data"]["systems"] = self.mixed_lmdb_path + # The fixture holds five 6-atom and five 9-atom frames; this budget + # puts a batch boundary inside the atom-count-sorted run, so at least + # one batch spans both sizes and is padded. + config["training"]["training_data"]["batch_size"] = "mix:27" + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + with tempfile.TemporaryDirectory(dir=self.tmpdir) as run_dir: + cwd = os.getcwd() + os.chdir(run_dir) + try: + trainer = get_trainer(config) + data = trainer.training_data + self.assertTrue(data._reader.mixed_nloc) + # Statistics keep their own fixed-nloc view, so padding never + # reaches the per-type accumulators. + self.assertEqual(data.get_stat_nsystems(), 2) + for sys_idx in range(data.get_stat_nsystems()): + stat = data.get_stat_batch(sys_idx) + self.assertTrue((stat["atype"] >= 0).all()) + padded = next( + batch + for batch in (data.get_batch() for _ in range(10)) + if (batch["atype"] < 0).any() + ) + self.assertEqual(padded["atype"].shape[1], 9) + self.assertEqual(padded["coord"].shape[1], 9) + self.assertEqual(padded["force"].shape[1], 9) + np.testing.assert_array_equal(padded["force"][padded["atype"] < 0], 0.0) + # Per-frame atom counts stay real, not padded. + self.assertEqual(sorted(set(padded["natoms"][:, 0].tolist())), [6, 9]) + trainer.run() + finally: + os.chdir(cwd) + + +class TestRaggedTrainingBatches(unittest.TestCase): + """A model reading a flat node axis is fed one, with nothing padded.""" + + def setUp(self) -> None: + self.tmpdir = tempfile.mkdtemp() + self.lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb") + _create_mixed_nloc_test_lmdb(self.lmdb_path) + + def tearDown(self) -> None: + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _config(self, descriptor: dict) -> dict: + config = { + "model": { + "type_map": ["O", "H"], + "descriptor": descriptor, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 1}, + "data_stat_nbatch": 1, + }, + "learning_rate": { + "type": "exp", + "decay_steps": 500, + "start_lr": 1e-3, + "stop_lr": 3.5e-8, + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + }, + "training": { + "training_data": { + "systems": self.lmdb_path, + "batch_size": "mix:27", + }, + "numb_steps": 2, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 100, + }, + } + return normalize(update_deepmd_input(config, warning=False)) + + @staticmethod + def _dpa1() -> dict: + return { + "type": "dpa1", + "sel": 12, + "rcut_smth": 0.5, + "rcut": 3.0, + "neuron": [8, 16], + "axis_neuron": 4, + "attn_layer": 0, + "precision": "float64", + "seed": 1, + } + + @staticmethod + def _se_e2_a() -> dict: + return { + "type": "se_e2_a", + "sel": [6, 12], + "rcut_smth": 0.5, + "rcut": 3.0, + "neuron": [8, 16], + "axis_neuron": 4, + "seed": 1, + } + + def _run(self, descriptor: dict, *, compile: bool = False): + """Train two steps and return the trainer plus one drawn batch.""" + config = self._config(descriptor) + config["training"]["enable_compile"] = compile + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + trainer = get_trainer(config) + batch = trainer.training_data.get_batch() + trainer.run() + return trainer, batch + finally: + os.chdir(cwd) + + def test_graph_model_trains_on_a_flat_node_axis(self) -> None: + """DPA1 reads a graph lower, so its batches carry no padded row.""" + trainer, batch = self._run(self._dpa1()) + self.assertTrue(trainer.training_data._reader.ragged_batches) + self.assertEqual(batch["coord"].ndim, 2) + self.assertEqual(batch["atype"].ndim, 1) + self.assertEqual(batch["coord"].shape[0], int(batch["n_node"].sum())) + self.assertTrue((batch["atype"] >= 0).all(), "nothing is padded") + # Frame-level fields keep their frame axis. + self.assertEqual(batch["energy"].shape[0], batch["n_node"].shape[0]) + + def test_compiled_graph_model_trains_on_a_flat_node_axis(self) -> None: + """The compiled lower reads the flat axis too, so compiling changes nothing. + + Its trace is taken on a rectangular system, which would bake in + ``N == nframes * nloc`` were the frame, node and edge counts not kept + as independent symbols. + """ + trainer, batch = self._run(self._dpa1(), compile=True) + self.assertEqual( + type(trainer.wrapper.model["Default"]).__name__, "_CompiledModel" + ) + self.assertTrue((batch["atype"] >= 0).all()) + self.assertEqual(batch["coord"].shape[0], int(batch["n_node"].sum())) + + def test_dense_model_keeps_padded_batches(self) -> None: + """se_e2_a reads a rectangular node axis and must still be padded.""" + trainer, batch = self._run(self._se_e2_a()) + self.assertFalse(trainer.training_data._reader.ragged_batches) + self.assertEqual(batch["coord"].ndim, 3) + self.assertEqual(batch["atype"].ndim, 2) + self.assertNotIn("n_node", batch) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/utils/test_nv_matrix_decode.py b/source/tests/pt_expt/utils/test_nv_matrix_decode.py index 5eccfa8277..43335230ff 100644 --- a/source/tests/pt_expt/utils/test_nv_matrix_decode.py +++ b/source/tests/pt_expt/utils/test_nv_matrix_decode.py @@ -1,16 +1,26 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""CPU unit tests for the nv dense-matrix -> (i, j, S) decode. +"""Parts of the nv path that run without a GPU. -The GPU ``neighbor_list`` search in ``build_neighbor_graph_nv`` is CUDA-only -and stays behind the opt-in CUDA suite (test_nv_graph_builder.py); the decode -(``nv_matrix_to_ijs``) is pure torch index arithmetic, so its regression-prone -parts (``// max_neighbors``, ``% nloc``, frame isolation, slot-validity mask) -are pinned here on the default CI with synthetic inputs. +Performance work on the nv builder is CUDA-bound and stays behind the opt-in +suite (test_nv_graph_builder.py), but two things are checkable anywhere +nvalchemiops imports. The decode (``nv_matrix_to_ijs``) is pure torch index +arithmetic, so its regression-prone parts (``// max_neighbors``, ``% nloc``, +frame isolation, slot-validity mask, and the lift back onto a padded batch) +are pinned here with synthetic inputs. The builder's agreement with the dense +reference builder is pinned here too, since the search itself is device +agnostic even though it is only worth running on a GPU. """ import numpy as np +import pytest import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) +from deepmd.pt.utils.nv_nlist import ( + is_nv_available, +) from deepmd.pt_expt.utils.nv_graph_builder import ( nv_matrix_to_ijs, ) @@ -66,6 +76,65 @@ def test_empty_no_neighbors(self) -> None: assert i.shape == (0,) and j.shape == (0,) assert s.shape == (0, 3) and f.shape == (0,) + def test_search_over_real_atoms_only(self) -> None: + """A search restricted to the real atoms still reports batch indices. + + A mixed-nloc batch is padded to a rectangular width, and the search is + given only the real slots, so its output is indexed on those. The decode + must lift both endpoints back onto the batch, or the frame and local + indices it derives would be meaningless. + + Batch: nf=2, nloc=3, real atoms at flat positions 0, 1, 3, 4, 5 -- that + is, frame 0 holds two atoms and frame 1 holds three. + """ + node_index = torch.tensor([0, 1, 3, 4, 5], dtype=torch.int64) + # Search indices 0..4 address those five atoms in order. + neighbor_matrix = torch.tensor( + [ + [1, 9], # search 0 (batch 0, frame 0 local 0) -> search 1 + [0, 9], # search 1 (batch 1, frame 0 local 1) -> search 0 + [3, 4], # search 2 (batch 3, frame 1 local 0) -> search 3, 4 + [2, 9], # search 3 (batch 4, frame 1 local 1) -> search 2 + [2, 9], # search 4 (batch 5, frame 1 local 2) -> search 2 + ], + dtype=torch.int32, + ) + num_neighbors = torch.tensor([1, 1, 2, 1, 1], dtype=torch.int32) + shifts = torch.zeros((5, 2, 3), dtype=torch.int32) + + i, j, s, f = nv_matrix_to_ijs( + neighbor_matrix, num_neighbors, shifts, 3, node_index=node_index + ) + assert _edge_set(i, j, s, f) == { + (0, 0, 1, (0, 0, 0)), + (0, 1, 0, (0, 0, 0)), + (1, 0, 1, (0, 0, 0)), + (1, 0, 2, (0, 0, 0)), + (1, 1, 0, (0, 0, 0)), + (1, 2, 0, (0, 0, 0)), + } + + def test_identity_node_index_matches_the_unrestricted_decode(self) -> None: + """Searching every slot is the case where the two indexings coincide.""" + rng = np.random.default_rng(3) + nf, nloc, mn = 2, 4, 3 + total = nf * nloc + num = torch.from_numpy(rng.integers(0, mn + 1, size=total)).to(torch.int32) + mat = torch.zeros((total, mn), dtype=torch.int32) + for dst in range(total): + frame = dst // nloc + mat[dst] = torch.from_numpy( + rng.integers(frame * nloc, (frame + 1) * nloc, size=mn) + ).to(torch.int32) + shf = torch.from_numpy(rng.integers(-1, 2, size=(total, mn, 3))).to(torch.int32) + + plain = nv_matrix_to_ijs(mat, num, shf, nloc) + lifted = nv_matrix_to_ijs( + mat, num, shf, nloc, node_index=torch.arange(total, dtype=torch.int64) + ) + for got, expected in zip(lifted, plain, strict=True): + torch.testing.assert_close(got, expected) + def test_random_vs_oracle(self) -> None: """Random matrices match a brute-force python oracle.""" rng = np.random.default_rng(11) @@ -97,3 +166,81 @@ def test_random_vs_oracle(self) -> None: nloc, ) assert _edge_set(i, j, s, f) == oracle + + +def _graph_edges(graph) -> set: + """Edges as (src, dst, rounded edge_vec), so two builders can be compared.""" + keep = graph.edge_mask + return { + (int(s), int(d), *(round(float(x), 8) for x in v)) + for s, d, v in zip( + graph.edge_index[0][keep], + graph.edge_index[1][keep], + graph.edge_vec[keep], + strict=True, + ) + } + + +@pytest.mark.skipif(not is_nv_available(), reason="nvalchemi-toolkit-ops not installed") +class TestNvBuilderOnPaddedBatches: + """The nv builder withholds phantom atoms from the search itself. + + Its cost grows with the square of the frame width, so a mixed-nloc batch + hands it the real atoms alone and lifts the resulting indices back onto the + padded batch. Both halves of that have to be exact. + """ + + @staticmethod + def _batch(nlocs, boxlen=9.0, seed=0): + rng = np.random.default_rng(seed) + width = max(nlocs) + coord = np.zeros((len(nlocs), width, 3)) + atype = np.full((len(nlocs), width), -1, dtype=np.int64) + for frame, nloc in enumerate(nlocs): + coord[frame, :nloc] = rng.uniform(0.0, boxlen, (nloc, 3)) + atype[frame, :nloc] = rng.integers(0, 2, nloc) + return ( + torch.tensor(coord), + torch.tensor(atype), + torch.tensor(np.tile(np.eye(3)[None] * boxlen, (len(nlocs), 1, 1))), + ) + + @pytest.mark.parametrize( + ("nlocs", "boxlen"), + [((6, 6, 6), 9.0), ((4, 7, 3), 9.0), ((5, 9, 2), 6.0)], + ids=["uniform", "mixed-nloc", "mixed-nloc-dense-box"], + ) + def test_matches_the_dense_builder(self, nlocs, boxlen) -> None: + from deepmd.pt_expt.utils.nv_graph_builder import ( + build_neighbor_graph_nv, + ) + + coord, atype, box = self._batch(nlocs, boxlen) + dense = build_neighbor_graph(coord, atype, box, 4.0) + nv = build_neighbor_graph_nv(coord, atype, box, 4.0) + assert _graph_edges(nv) == _graph_edges(dense) + np.testing.assert_array_equal(nv.n_node.numpy(), dense.n_node.numpy()) + + def test_search_enumerates_the_real_atoms_only(self, monkeypatch) -> None: + """Padding must not reach the search, only its output indexing.""" + from deepmd.pt_expt.utils import ( + nv_graph_builder, + ) + + nlocs = (4, 7, 3) + coord, atype, box = self._batch(nlocs) + searched: list[int] = [] + original = nv_graph_builder.nv_search_matrix + + def spy(*args, **kwargs): + result = original(*args, **kwargs) + searched.append(int(result[2].shape[0])) + return result + + monkeypatch.setattr(nv_graph_builder, "nv_search_matrix", spy) + nv_graph_builder.build_neighbor_graph_nv(coord, atype, box, 4.0) + assert searched == [sum(nlocs)], ( + f"the search saw {searched} atoms; the batch holds {sum(nlocs)} real " + f"atoms padded to {len(nlocs) * max(nlocs)} slots" + ) From a809f15fe04c31d7560b92b432dba3b030b8a370 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 31 Jul 2026 17:49:24 +0800 Subject: [PATCH 2/9] fix(lmdb): make label availability lazy and requirement-aware Large LMDBs must not pay O(frame count) random I/O or Python-object allocation before the active training contract is known. Read metadata through sequential readahead, keep frame tables in compact NumPy arrays, and choose readahead according to the access pattern of each reader. Defer availability resolution until requirements are registered. Probe only optional tracked fields; uniform datasets start without a full scan, while detected mixed datasets build one compact cached signature index through a sequential reader with bounded progress logging. Mandatory fields fail at decode, default-backed inputs remain available per frame, and derived fields are computed from normalized structure data. Apply the contract consistently to statistics, samplers, full validation, and both PT training paths. Declare only active loss labels, preserve explicit values beside defaults, and gate force-derived losses by the availability of their force target. Keep filtered frame and system indices, mixed-nloc packing, and validation views in one index domain, and retain the block-allocation and batch-layout fixes found while consolidating the data path. --- deepmd/dpmodel/loss/ener.py | 201 +- deepmd/dpmodel/loss/ener_spin.py | 1 + deepmd/dpmodel/loss/tensor.py | 1 + deepmd/dpmodel/utils/lmdb_data.py | 1750 ++++++++++++----- deepmd/pt/loss/ener.py | 22 +- deepmd/pt/loss/tensor.py | 1 + deepmd/pt/train/training.py | 5 + deepmd/pt/utils/lmdb_dataset.py | 85 +- deepmd/pt_expt/train/training.py | 15 +- deepmd/pt_expt/train/validation.py | 17 +- deepmd/pt_expt/utils/lmdb_dataset.py | 8 +- deepmd/utils/data.py | 23 + source/tests/common/dpmodel/test_lmdb_data.py | 666 ++++++- source/tests/common/dpmodel/test_loss_ener.py | 97 + source/tests/pt/test_lmdb_dataloader.py | 80 +- source/tests/pt/test_loss_default_pf.py | 73 + .../pt_expt/model/test_dpa4_native_spin.py | 1 + source/tests/pt_expt/test_lmdb_training.py | 177 ++ source/tests/pt_expt/test_training.py | 2 + 19 files changed, 2516 insertions(+), 709 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 46fc67af51..1dd70361b7 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -253,32 +253,31 @@ def call( including all enabled prefactors and any configured Huber terms. """ energy = model_dict["energy"] - force = model_dict["force"] - virial = model_dict["virial"] - atom_ener = model_dict["atom_energy"] - energy_hat = label_dict["energy"] - force_hat = label_dict["force"] - virial_hat = label_dict["virial"] - atom_ener_hat = label_dict["atom_ener"] - atom_pref = label_dict["atom_pref"] - find_energy = label_dict["find_energy"] - find_force = label_dict["find_force"] - find_virial = label_dict["find_virial"] - find_atom_ener = label_dict["find_atom_ener"] - find_atom_pref = ( - label_dict["find_atom_pref"] if not self.use_default_pf else 1.0 - ) - xp = array_api_compat.array_namespace( - energy, - force, - virial, - atom_ener, - energy_hat, - force_hat, - virial_hat, - atom_ener_hat, - atom_pref, + xp = array_api_compat.array_namespace(energy) + + force_required = ( + self.has_f or self.has_pf or self.relative_f is not None or self.has_gf ) + if self.has_e: + energy_hat = label_dict["energy"] + find_energy = label_dict["find_energy"] + if force_required: + force = model_dict["force"] + force_hat = label_dict["force"] + find_force = label_dict["find_force"] + if self.has_v: + virial = model_dict["virial"] + virial_hat = label_dict["virial"] + find_virial = label_dict["find_virial"] + if self.has_ae: + atom_ener = model_dict["atom_energy"] + atom_ener_hat = label_dict["atom_ener"] + find_atom_ener = label_dict["find_atom_ener"] + if self.has_pf: + atom_pref = label_dict["atom_pref"] + find_atom_pref = ( + label_dict["find_atom_pref"] if not self.use_default_pf else 1.0 + ) # Two things about a batch decide how its terms reduce, and the node # axis states them differently. @@ -309,10 +308,11 @@ def call( # A + B -> C + D # E = - E(A) - E(B) + E(C) + E(D) # A, B, C, D could be put far away from each other + atom_ener = model_dict["atom_energy"] atom_ener_coeff = label_dict["atom_ener_coeff"] atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener.shape) energy = xp.sum(atom_ener_coeff * atom_ener, axis=1) - if self.has_f or self.has_pf or self.relative_f or self.has_gf: + if force_required: force_reshape = xp.reshape(force, (-1,)) force_hat_reshape = xp.reshape(force_hat, (-1,)) diff_f = force_hat_reshape - force_reshape @@ -332,22 +332,33 @@ def call( atom_norm = 1.0 / natoms atom_norm_ener = 1.0 / natoms lr_ratio = learning_rate / self.starter_learning_rate - pref_e = find_energy * ( - self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * lr_ratio - ) - pref_f = find_force * ( - self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * lr_ratio - ) - pref_v = find_virial * ( - self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * lr_ratio - ) - pref_ae = find_atom_ener * ( - self.limit_pref_ae + (self.start_pref_ae - self.limit_pref_ae) * lr_ratio - ) - pref_pf = find_atom_pref * ( - self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio - ) - pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio + if self.has_e: + pref_e = find_energy * ( + self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * lr_ratio + ) + if self.has_f: + pref_f = find_force * ( + self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * lr_ratio + ) + if self.has_v: + pref_v = find_virial * ( + self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * lr_ratio + ) + if self.has_ae: + pref_ae = find_atom_ener * ( + self.limit_pref_ae + + (self.start_pref_ae - self.limit_pref_ae) * lr_ratio + ) + if self.has_pf: + effective_find_pf = find_force * find_atom_pref + pref_pf = effective_find_pf * ( + self.limit_pref_pf + + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio + ) + if self.has_h: + pref_h = ( + self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio + ) loss = 0 more_loss = {} @@ -693,12 +704,12 @@ def call( ) loss += pref_pf * l2_pf_masked more_loss["rmse_pf"] = self.display_if_exist( - xp.sqrt(l2_pf_masked), find_atom_pref + xp.sqrt(l2_pf_masked), effective_find_pf ) else: loss += pref_pf * l2_pref_force_loss more_loss["rmse_pf"] = self.display_if_exist( - xp.sqrt(l2_pref_force_loss), find_atom_pref + xp.sqrt(l2_pref_force_loss), effective_find_pf ) elif self.loss_func == "mae": l1_pref_force_loss = xp.mean( @@ -710,12 +721,12 @@ def call( l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3) loss += pref_pf * l1_pf_masked more_loss["mae_pf"] = self.display_if_exist( - l1_pf_masked, find_atom_pref + l1_pf_masked, effective_find_pf ) else: loss += pref_pf * l1_pref_force_loss more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss, find_atom_pref + l1_pref_force_loss, effective_find_pf ) else: raise NotImplementedError( @@ -735,7 +746,8 @@ def call( ) find_drdq = label_dict["find_drdq"] drdq = label_dict["drdq"] - pref_gf = find_drdq * ( + effective_find_gf = find_force * find_drdq + pref_gf = effective_find_gf * ( self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * lr_ratio ) @@ -771,7 +783,7 @@ def call( l2_gen_force_loss = xp.mean(xp.square(diff_gen_force)) loss += pref_gf * l2_gen_force_loss more_loss["rmse_gf"] = self.display_if_exist( - xp.sqrt(l2_gen_force_loss), find_drdq + xp.sqrt(l2_gen_force_loss), effective_find_gf ) hessian = model_dict.get("hessian", model_dict.get("energy_derv_r_derv_r")) if self.has_h and hessian is not None and "hessian" in label_dict: @@ -818,54 +830,60 @@ def call( @property def label_requirement(self) -> list[DataRequirementItem]: """Return data label requirements needed for this loss calculation.""" - label_requirement = [] - label_requirement.append( - DataRequirementItem( - "energy", - ndof=1, - atomic=False, - must=False, - high_prec=True, + label_requirement: list[DataRequirementItem] = [] + if self.has_e: + label_requirement.append( + DataRequirementItem( + "energy", + ndof=1, + atomic=False, + must=False, + high_prec=True, + ) ) - ) - label_requirement.append( - DataRequirementItem( - "force", - ndof=3, - atomic=True, - must=False, - high_prec=False, + if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf: + label_requirement.append( + DataRequirementItem( + "force", + ndof=3, + atomic=True, + must=False, + high_prec=False, + ) ) - ) - label_requirement.append( - DataRequirementItem( - "virial", - ndof=9, - atomic=False, - must=False, - high_prec=False, + if self.has_v: + label_requirement.append( + DataRequirementItem( + "virial", + ndof=9, + atomic=False, + must=False, + high_prec=False, + ) ) - ) - label_requirement.append( - DataRequirementItem( - "atom_ener", - ndof=1, - atomic=True, - must=False, - high_prec=False, + if self.has_ae: + label_requirement.append( + DataRequirementItem( + "atom_ener", + ndof=1, + atomic=True, + must=False, + high_prec=False, + ) ) - ) - label_requirement.append( - DataRequirementItem( - "atom_pref", - ndof=1, - atomic=True, - must=False, - high_prec=False, - repeat=3, - default=1.0, + if self.has_pf: + label_requirement.append( + DataRequirementItem( + "atom_pref", + ndof=1, + atomic=True, + must=False, + high_prec=False, + repeat=3, + default=1.0, + source_policy="default" if self.use_default_pf else "tracked", + ) ) - ) if self.has_gf > 0: label_requirement.append( DataRequirementItem( @@ -885,6 +903,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, default=1.0, + source_policy="default", ) ) if self.has_h: diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 30959fa83f..4926c64813 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -471,6 +471,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, default=1.0, + source_policy="default", ) ) return label_requirement diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index 11a8ce3987..9af03a3b85 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -197,6 +197,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, default=1.0, + source_policy="default", ) ) return label_requirement diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 90cb0a5ae1..888daf3166 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -9,9 +9,13 @@ import logging import math import multiprocessing +import os import signal import threading +import time from collections.abc import ( + Callable, + Iterable, Iterator, Sequence, ) @@ -26,14 +30,12 @@ from dataclasses import ( dataclass, ) -from itertools import ( - pairwise, -) from pathlib import ( Path, ) from typing import ( Any, + cast, ) import lmdb @@ -50,10 +52,17 @@ from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, + DataRequirementSourcePolicy, ) log = logging.getLogger(__name__) + +def _is_local_rank_zero() -> bool: + """Whether this process owns node-local operational logging.""" + return int(os.environ.get("LOCAL_RANK", "0")) == 0 + + # LMDB key → DeePMD convention _KEY_REMAP = { "coords": "coord", @@ -68,20 +77,26 @@ # (energy is set by Loss DataRequirementItem; reduce() also sets high_prec=True) _HIGH_PREC_KEYS = frozenset({"energy"}) -# Keys that describe frame geometry or LMDB bookkeeping rather than optional -# model inputs/labels. They must not participate in availability signatures. -_STRUCTURAL_KEYS = frozenset( - { - "coord", - "box", - "atype", - "natoms", - "real_natoms_vec", - "fid", - } -) -_LMDB_METADATA_KEYS = frozenset({"atom_numbs", "atom_names", "orig"}) -_OPTIONAL_MODEL_INPUT_KEYS = frozenset({"fparam", "aparam", "spin", "charge_spin"}) +# Frames probed to decide whether availability-sensitive labels partition a +# dataset. Requirements are registered after reader construction, so the probe +# is deferred until a consumer first requests sampling groups. +# +# The sample makes the partition probabilistic, and a dataset whose odd frame +# out falls between two probes is grouped as if it were uniform. What such a +# frame costs is bounded: the batch it lands in reports the disputed label +# unavailable, so at most ``batch_size - 1`` frames lose that label for that +# batch, and a reshuffle moves it elsewhere next epoch. No default fill is +# ever mistaken for a real label, which is the property that matters -- see +# :func:`_batch_find_flags`. +# +# Raising this bound trades startup reads for a tighter guarantee, but never +# reaches certainty. A dataset that really does interleave labels is better +# fixed where it is written: recording a per-frame availability id in +# ``__metadata__``, alongside the ``frame_nlocs`` and ``frame_system_ids`` +# already there, would make the partition exact and free. +_AVAILABILITY_PROBE_FRAMES = 256 +_AVAILABILITY_SCAN_CHUNK = 4096 +_AVAILABILITY_LOG_SECONDS = 60.0 # Atom type written into the padded slots of a mixed-nloc batch. A phantom # atom occupies a tensor slot but no physical site: the neighbor list gives it @@ -121,7 +136,7 @@ _ENV_CACHE: dict[str, tuple[lmdb.Environment, int]] = {} -def _open_lmdb(path: str) -> lmdb.Environment: +def _open_lmdb(path: str, *, sequential: bool = False) -> lmdb.Environment: """Open (or reuse) a readonly LMDB environment with reference counting. The python-lmdb binding raises ``lmdb.Error`` if the same path is opened @@ -129,6 +144,27 @@ def _open_lmdb(path: str) -> lmdb.Environment: and bump a reference count. Call :func:`_close_lmdb` when done to decrement the count; when it reaches zero the environment is closed and removed from the cache. + + Parameters + ---------- + path : str + Path to the LMDB directory. + + sequential : bool, optional + Whether the caller reads frames in ascending key order. Kernel + readahead is then left on, which turns the many small faults such a + walk would make into few large reads. Measured against NFS: 45 000 + frames/s with readahead against 1 300 without. A caller reading in + shuffled order leaves this false, where readahead instead fetches + neighbours it will not use and costs about 12% of the throughput. + + One path admits one environment, so a path already open keeps the + setting its first caller asked for. + + Returns + ------- + lmdb.Environment + The shared read-only environment for ``path``. """ resolved = str(Path(path).resolve()) entry = _ENV_CACHE.get(resolved) @@ -136,7 +172,9 @@ def _open_lmdb(path: str) -> lmdb.Environment: env, refcount = entry _ENV_CACHE[resolved] = (env, refcount + 1) return env - env = lmdb.open(path, readonly=True, lock=False, readahead=False, meminit=False) + env = lmdb.open( + path, readonly=True, lock=False, readahead=sequential, meminit=False + ) _ENV_CACHE[resolved] = (env, 1) return env @@ -166,6 +204,43 @@ def _read_metadata(txn: lmdb.Transaction) -> dict: return msgpack.unpackb(raw, raw=False) +def _read_metadata_of(path: str) -> dict: + """Read ``__metadata__`` through an environment that keeps readahead on. + + The value is a single contiguous run of overflow pages -- hundreds of + megabytes once the writer records a per-frame table. ``MDB_NORDAHEAD``, + which a reader of shuffled frames asks :func:`_open_lmdb` for, advises + the kernel ``MADV_RANDOM`` over the whole map, and that turns this + sequential run into one synchronous fault per 4 KiB page. Measured + against NFS on a 665 MiB table: 379 s without readahead, 8.2 s with it. + + The environment is closed before the caller opens its own, because + python-lmdb refuses a second open of one path. A path already open for + frame serving is read through that environment instead, its pages being + warm by then in the only case that matters. + + Parameters + ---------- + path : str + Path to the LMDB directory. + + Returns + ------- + dict + The decoded metadata mapping. + """ + entry = _ENV_CACHE.get(str(Path(path).resolve())) + if entry is not None: + with entry[0].begin() as transaction: + return _read_metadata(transaction) + env = lmdb.open(path, readonly=True, lock=False, readahead=True, meminit=False) + try: + with env.begin() as transaction: + return _read_metadata(transaction) + finally: + env.close() + + def _decode_array(obj: dict, *, copy: bool = True) -> np.ndarray: """Reconstruct ndarray from msgpack-encoded dict with {type, shape, data}. @@ -226,23 +301,358 @@ def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]: return out -def _availability_signature_keys( - frame: dict[str, Any], requirement_keys: Iterator[str] -) -> list[str]: - """Return data keys whose availability can affect frame collation. +def _requirement_value(requirement: Any, key: str, default: Any) -> Any: + """Read one requirement attribute from object or dictionary form.""" + if isinstance(requirement, dict): + return requirement.get(key, default) + return getattr(requirement, key, default) + + +def _requirement_source_policy(requirement: Any) -> DataRequirementSourcePolicy: + """Return the normalized source-presence policy of one requirement.""" + policy = str(_requirement_value(requirement, "source_policy", "tracked")) + if policy not in {"tracked", "default", "derived"}: + raise ValueError(f"Unsupported data requirement source policy {policy!r}") + return cast("DataRequirementSourcePolicy", policy) + + +def _requirement_is_mandatory(requirement: Any) -> bool: + """Whether unavailable source data violates one requirement.""" + return bool(_requirement_value(requirement, "must", False)) + + +def _availability_requirement_keys( + requirements: dict[str, Any], +) -> tuple[str, ...]: + """Return requirements whose source presence affects consumer behavior. + + Only optional ``tracked`` fields participate. Mandatory fields fail at + decode time rather than forming a default-filled group. ``default`` and + ``derived`` fields have valid per-frame fallbacks and therefore combine + safely regardless of source presence. + + Parameters + ---------- + requirements : dict[str, Any] + Data requirements keyed by normalized field name. + + Returns + ------- + tuple[str, ...] + Sorted availability-sensitive field names. + """ + return tuple( + sorted( + key + for key, requirement in requirements.items() + if _requirement_source_policy(requirement) == "tracked" + and not _requirement_is_mandatory(requirement) + ) + ) + + +def _raw_frame_availability( + raw: bytes, + key_bits: dict[str, int], +) -> int: + """Return one undecoded frame's selected availability bit mask. + + A selected field is available when the frame carries it unless an + explicit ``find_*`` flag says otherwise. Fields outside + ``key_bits`` are deliberately ignored, so auxiliary raw data cannot alter + a training run's batch partition. + + Only the msgpack map is walked; the arrays it describes stay encoded. + + Parameters + ---------- + raw : bytes + The msgpack payload of one LMDB frame. + key_bits : dict[str, int] + Bit assigned to each normalized availability-sensitive field. + + Returns + ------- + int + Availability mask in ``key_bits`` bit positions. + """ + frame = msgpack.unpackb(raw, raw=False) + present = 0 + explicit_known = 0 + explicit_true = 0 + for key, value in frame.items(): + name = _KEY_REMAP.get(key, key) + if name.startswith("find_"): + label = name.removeprefix("find_") + bit = key_bits.get(label) + if bit is not None: + explicit_known |= bit + if float(np.asarray(_decode_value(value)).item()) != 0.0: + explicit_true |= bit + else: + bit = key_bits.get(name) + if bit is not None: + present |= bit + return present & (~explicit_known | explicit_true) + + +def _evenly_spaced(keys: Sequence[int], count: int) -> list[int]: + """Pick up to ``count`` entries spread evenly over ``keys``. + + Spreading rather than truncating matters for a dataset assembled by + concatenating sources: the tail of the key space is as likely to hold + the odd one out as the head. + + Parameters + ---------- + keys : Sequence[int] + Integer LMDB frame keys to sample from. + count : int + Upper bound on the number of keys returned. + + Returns + ------- + list[int] + The sampled keys, in ascending position order. + """ + total = len(keys) + if total <= count: + return [int(key) for key in keys] + return [int(keys[total * position // count]) for position in range(count)] + + +def _probe_uniform_availability( + transaction: lmdb.Transaction, + keys: Iterable[int], + frame_format: str, + availability_keys: Sequence[str], +) -> bool: + """Whether the sampled frames all supply the same labels. + + Parameters + ---------- + transaction : lmdb.Transaction + Open read transaction on the LMDB environment. + keys : Iterable[int] + Integer LMDB frame keys to probe, already reduced to a bounded + sample by :func:`_evenly_spaced`. + frame_format : str + Format specification for integer LMDB frame keys. + availability_keys : Sequence[str] + Normalized field names used by the active consumer. + + Returns + ------- + bool + True when every probed frame supplies the same labels. A dataset + larger than the sample may still be mixed, in which case the batch + decode reports the disputed label unavailable rather than mixing it. + """ + key_bits = {key: 1 << position for position, key in enumerate(availability_keys)} + reference: int | None = None + for key in keys: + raw = transaction.get(format(int(key), frame_format).encode()) + if raw is None: + continue + availability = _raw_frame_availability(raw, key_bits) + if reference is None: + reference = availability + elif availability != reference: + return False + return True + + +@dataclass(frozen=True) +class _AvailabilityIndex: + """Compact availability signature ID aligned with a frame index domain.""" + + ids: np.ndarray + signature_count: int - In addition to registered requirements and standard optional model inputs, - include every label-like field that :class:`LmdbDataReader` exposes from - the raw frame. This keeps sampler/validation grouping consistent with the - complete set of ``find_*`` flags checked during collation. + def groups( + self, + indices: np.ndarray, + *, + positions: np.ndarray | None = None, + ) -> list[np.ndarray]: + """Partition indices by cached signature IDs without source reads.""" + index_array = np.asarray(indices) + signature_ids = ( + self.ids + if positions is None + else self.ids[np.asarray(positions, dtype=np.int64)] + ) + if len(index_array) != len(signature_ids): + raise ValueError( + "availability index and frame indices have different lengths: " + f"{len(signature_ids)} != {len(index_array)}" + ) + if len(index_array) == 0: + return [] + if self.signature_count == 1: + return [index_array] + + if self.signature_count <= 8: + groups: list[np.ndarray] = [] + for signature_id in range(self.signature_count): + mask = signature_ids == signature_id + if np.any(mask): + groups.append(index_array[mask]) + return groups + + order = np.argsort(signature_ids, kind="stable") + ordered_ids = signature_ids[order] + cuts = np.flatnonzero(ordered_ids[1:] != ordered_ids[:-1]) + 1 + ordered_indices = index_array[order] + return list(np.split(ordered_indices, cuts)) + + +def _widen_signature_ids(ids: np.ndarray, written: int) -> np.ndarray: + """Return the next wider unsigned ID buffer, preserving written entries.""" + widths = { + np.dtype(np.uint8): np.dtype(np.uint16), + np.dtype(np.uint16): np.dtype(np.uint32), + np.dtype(np.uint32): np.dtype(np.uint64), + } + target_dtype = widths.get(ids.dtype) + if target_dtype is None: + raise OverflowError("LMDB availability signatures exceed uint64 capacity") + widened = np.empty(ids.shape, dtype=target_dtype) + widened[:written] = ids[:written] + return widened + + +def _scan_availability_index( + frame_count: int, + read_raw: Callable[[int], bytes | None], + availability_keys: Sequence[str], + dataset: str, +) -> _AvailabilityIndex: + """Build compact signature IDs with one exact source scan. + + The scan stores one unsigned integer per frame and one dictionary entry per + distinct signature. It never accumulates frame indices as Python objects. + + Parameters + ---------- + frame_count : int + Number of positions in the index domain. + read_raw : Callable[[int], bytes or None] + Function returning the encoded frame for one domain position. + availability_keys : Sequence[str] + Normalized field names defining the partition. + dataset : str + Dataset path included in progress messages. + + Returns + ------- + _AvailabilityIndex + Compact signature IDs aligned with domain positions. """ - keys = set(requirement_keys) | set(_OPTIONAL_MODEL_INPUT_KEYS) - for frame_key in frame: - if frame_key.startswith("find_"): - keys.add(frame_key.removeprefix("find_")) - elif frame_key not in _STRUCTURAL_KEYS | _LMDB_METADATA_KEYS: - keys.add(frame_key) - return sorted(keys) + report_progress = _is_local_rank_zero() + if report_progress: + log.info( + "LMDB label-availability scan started: dataset=%s, frames=%d, labels=%s", + dataset, + frame_count, + list(availability_keys), + ) + + key_bits = {key: 1 << position for position, key in enumerate(availability_keys)} + signature_ids = np.empty(frame_count, dtype=np.uint8) + signature_map: dict[int, int] = {} + next_log = time.monotonic() + _AVAILABILITY_LOG_SECONDS if report_progress else 0.0 + for start in range(0, frame_count, _AVAILABILITY_SCAN_CHUNK): + stop = min(start + _AVAILABILITY_SCAN_CHUNK, frame_count) + for position in range(start, stop): + raw = read_raw(position) + if raw is None: + raise RuntimeError( + f"LMDB frame at position {position} is missing from {dataset}" + ) + signature = _raw_frame_availability(raw, key_bits) + signature_id = signature_map.get(signature) + if signature_id is None: + signature_id = len(signature_map) + signature_map[signature] = signature_id + if signature_id > np.iinfo(signature_ids.dtype).max: + signature_ids = _widen_signature_ids(signature_ids, position) + signature_ids[position] = signature_id + + if report_progress: + now = time.monotonic() + if stop < frame_count and now >= next_log: + log.info( + "LMDB label-availability scan progress: dataset=%s, " + "frames=%d/%d (%.1f%%)", + dataset, + stop, + frame_count, + 100.0 * stop / frame_count, + ) + next_log = now + _AVAILABILITY_LOG_SECONDS + + if report_progress: + log.info( + "LMDB label-availability scan completed: dataset=%s, frames=%d, " + "groups=%d, index_dtype=%s", + dataset, + frame_count, + len(signature_map), + signature_ids.dtype, + ) + return _AvailabilityIndex(signature_ids, len(signature_map)) + + +def _scan_lmdb_path_sequential( + lmdb_path: str, + availability_keys: Sequence[str], + log_level: int, +) -> _AvailabilityIndex: + """Scan a complete LMDB under a sequential-readahead environment.""" + logging.basicConfig(level=log_level) + logging.getLogger().setLevel(log_level) + environment = lmdb.open( + lmdb_path, + readonly=True, + lock=False, + readahead=True, + meminit=False, + ) + try: + with environment.begin() as transaction: + metadata = _read_metadata(transaction) + frame_count, frame_format, _natoms_per_type = _parse_metadata(metadata) + with environment.begin() as transaction: + + def read_raw(position: int) -> bytes | None: + return transaction.get(format(position, frame_format).encode()) + + return _scan_availability_index( + frame_count, + read_raw, + availability_keys, + lmdb_path, + ) + finally: + environment.close() + + +def _scan_lmdb_path_in_worker( + lmdb_path: str, + availability_keys: Sequence[str], +) -> _AvailabilityIndex: + """Run a sequential scan outside a process holding this LMDB open.""" + executor = _create_lmdb_executor(1) + try: + return executor.submit( + _scan_lmdb_path_sequential, + lmdb_path, + availability_keys, + log.getEffectiveLevel(), + ).result() + finally: + executor.shutdown(wait=True, cancel_futures=True) def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray: @@ -276,12 +686,15 @@ class LmdbDecodeConfig: Optional LMDB-type to model-type lookup table. data_requirements Registered data requirements keyed by field name. + dataset + Dataset identifier used in frame-level diagnostics. """ ntypes: int natoms: int type_remap: np.ndarray | None data_requirements: dict[str, Any] + dataset: str = "" def _requirement_dtype(requirement: Any) -> np.dtype: @@ -316,6 +729,48 @@ def _requirement_is_atomic(requirement: Any) -> bool: return bool(getattr(requirement, "atomic", False)) +def _frame_source_available(frame: dict[str, Any], key: str) -> bool: + """Resolve one frame's explicit or inferred source availability.""" + value = frame.get(key) + source_present = _is_encoded_array(value) or isinstance( + value, (np.ndarray, np.generic, int, float, bool) + ) + find_key = f"find_{key}" + if find_key in frame: + return source_present and bool( + float(np.asarray(_decode_value(frame[find_key])).item()) + ) + return source_present + + +def _raise_if_mandatory_unavailable( + frame: dict[str, Any], + key: str, + requirement: Any, + source_available: bool, + *, + dataset: str, + frame_index: int, +) -> None: + """Reject unavailable mandatory data at the shared frame boundary.""" + if not _requirement_is_mandatory(requirement) or source_available: + return + find_key = f"find_{key}" + if find_key in frame: + find_value = float(np.asarray(_decode_value(frame[find_key])).item()) + reason = ( + f"explicit {find_key}=0" + if find_value == 0.0 + else f"field is absent or invalid despite {find_key}={find_value:g}" + ) + else: + reason = "field is absent or invalid" + raise RuntimeError( + f"Required LMDB field {key!r} is unavailable in frame {frame_index} " + f"of {dataset}: {reason}." + ) + + def resolve_per_atom_keys( frame: dict[str, Any], config: LmdbDecodeConfig, @@ -383,6 +838,34 @@ def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray: return natoms +def _resolve_derived_requirement( + frame: dict[str, Any], + key: str, + requirement: Any, + config: LmdbDecodeConfig, +) -> None: + """Resolve one derived field from normalized structural frame data.""" + if key != "min_pair_dist": + raise ValueError(f"Unsupported derived LMDB field {key!r}") + + coord = frame.get("coord") + atype = frame.get("atype") + if not isinstance(coord, np.ndarray) or not isinstance(atype, np.ndarray): + frame.pop(key, None) + frame[f"find_{key}"] = np.float32(0.0) + return + + box = frame.get("box") + if box is not None and np.allclose(box, 0.0): + box = None + threshold = float(_requirement_value(requirement, "default", 0.0)) + frame[key] = np.array( + [compute_min_pair_dist_single(coord, box, atype, stop_below=threshold)], + dtype=_resolve_frame_dtype(config, key), + ) + frame[f"find_{key}"] = np.float32(1.0) + + def decode_lmdb_frame( raw: bytes, original_key: int, @@ -468,34 +951,9 @@ def decode_lmdb_frame( frame["real_natoms_vec"] = natoms requirements = config.data_requirements - coord = frame.get("coord") - if ( - "min_pair_dist" in requirements - and "min_pair_dist" not in frame - and isinstance(coord, np.ndarray) - and isinstance(atype, np.ndarray) - ): - box = frame.get("box") - if box is not None and np.allclose(box, 0.0): - box = None - requirement = requirements["min_pair_dist"] - default = ( - requirement.get("default", 0.0) - if isinstance(requirement, dict) - else getattr(requirement, "default", 0.0) - ) - frame["find_min_pair_dist"] = np.float32(1.0) - frame["min_pair_dist"] = np.array( - [ - compute_min_pair_dist_single( - coord, - box, - atype, - stop_below=float(default), - ) - ], - dtype=_resolve_frame_dtype(config, "min_pair_dist"), - ) + for key, requirement in requirements.items(): + if _requirement_source_policy(requirement) == "derived": + _resolve_derived_requirement(frame, key, requirement, config) structural_keys = frozenset( { @@ -524,16 +982,28 @@ def decode_lmdb_frame( atomic = requirement.atomic repeat = getattr(requirement, "repeat", 1) dtype = _requirement_dtype(requirement) + find_key = f"find_{key}" + source_available = _frame_source_available(frame, key) + _raise_if_mandatory_unavailable( + frame, + key, + requirement, + source_available, + dataset=config.dataset, + frame_index=original_key, + ) - if key not in frame: - frame[f"find_{key}"] = np.float32(0.0) + if not source_available: + frame[find_key] = np.float32( + 1.0 if _requirement_source_policy(requirement) == "default" else 0.0 + ) shape = (frame_natoms, ndof) if atomic else (ndof,) data = np.full(shape, default, dtype=dtype) if repeat != 1: data = np.repeat(data, repeat).reshape(-1) frame[key] = data else: - frame.setdefault(f"find_{key}", np.float32(1.0)) + frame.setdefault(find_key, np.float32(1.0)) if repeat != 1 and isinstance(frame[key], np.ndarray): frame[key] = ( np.repeat(frame[key], repeat).reshape(-1).astype(dtype, copy=False) @@ -789,13 +1259,26 @@ def decode_lmdb_batch( The layout fixes the shape of every field of the result, so a chunked decode must pass each chunk the layout of its own frames, cut from the batch-wide one; :meth:`LmdbDataReader.batch_layout` resolves that once. + + Frames need not expose the same optional fields. A field only some of + them carry is left out of the batch, and the matching ``find_*`` flag + reports the label unavailable, which matches what + :func:`collate_lmdb_frames` produces for the same frames. Registered + data requirements are exempt: :func:`decode_lmdb_frame` puts them on + every frame. """ if not original_keys: raise ValueError("decode_lmdb_batch requires at least one frame key") batch: dict[str, Any] | None = None batch_size = len(original_keys) - expected_fields: frozenset[str] | None = None + # Frames need not expose the same optional fields. Counting the frames + # each field appears on settles, once the batch is read, which fields + # have a complete column to stack and which labels the batch may report + # as available; see :func:`_batch_find_flags` for the same rule applied + # to the generic collation path. + field_counts: dict[str, int] = {} + find_flags: dict[str, bool] = {} for row, original_key in enumerate(original_keys): key = format(int(original_key), frame_format).encode() raw = transaction.get(key) @@ -819,27 +1302,23 @@ def decode_lmdb_batch( f"the batch layout gives frame {original_key} " f"{int(layout.n_node[row])} atoms, but it holds {frame_nloc}" ) + for field in frame: + field_counts[field] = field_counts.get(field, 0) + 1 + if batch is None: batch = _allocate_lmdb_batch(frame, batch_size, layout) - expected_fields = frozenset(frame) + find_flags = { + field: float(value) != 0.0 + for field, value in frame.items() + if field.startswith("find_") + } continue - frame_fields = frozenset(frame) - if frame_fields != expected_fields: - raise ValueError( - "LMDB frames in one batch expose inconsistent fields: " - f"frame {original_keys[0]} has {sorted(expected_fields)}, while " - f"frame {original_key} has {sorted(frame_fields)}" - ) for field, value in frame.items(): if field.startswith("find_"): - if not np.array_equal(batch[field], value): - raise ValueError( - f"LMDB field availability changes within one batch: " - f"{field!r} differs at frame {original_key}" - ) + find_flags[field] = find_flags.get(field, False) and float(value) != 0.0 continue - if field == "type" or value is None: + if field == "type" or value is None or field not in batch: continue if field == "fid": batch[field][row] = value @@ -880,6 +1359,12 @@ def decode_lmdb_batch( destination[row] = array assert batch is not None and layout is not None + for field, present in find_flags.items(): + available = present and field_counts.get(field, 0) == batch_size + batch[field] = np.float32(1.0 if available else 0.0) + for field, count in field_counts.items(): + if count < batch_size and not field.startswith("find_"): + batch.pop(field, None) if layout.ragged: batch["n_node"] = layout.n_node batch["sid"] = np.asarray([0], dtype=np.int64) @@ -926,39 +1411,38 @@ def _decode_lmdb_worker_chunk( def _merge_lmdb_chunks(chunks: list[dict[str, Any]]) -> dict[str, Any]: - """Merge ordered worker chunks into one contiguous batch.""" + """Merge ordered worker chunks into one contiguous batch. + + Each chunk has already reduced its own frames, so merging repeats that + reduction one level up: a label is available to the merged batch only + where every chunk reports it, and a field only some chunks carry has no + complete column to concatenate. + """ if not chunks: raise ValueError("cannot merge an empty LMDB chunk list") if len(chunks) == 1: return chunks[0] first = chunks[0] - expected_fields = frozenset(first) - for chunk_index, chunk in enumerate(chunks[1:], start=1): - chunk_fields = frozenset(chunk) - if chunk_fields != expected_fields: - raise ValueError( - "LMDB worker chunks expose inconsistent fields: " - f"chunk 0 has {sorted(expected_fields)}, while chunk " - f"{chunk_index} has {sorted(chunk_fields)}" - ) - merged: dict[str, Any] = {} for key, value in first.items(): if key.startswith("find_"): - for chunk_index, chunk in enumerate(chunks[1:], start=1): - if not np.array_equal(value, chunk[key]): - raise ValueError( - "LMDB field availability changes across worker chunks: " - f"{key!r} differs in chunk {chunk_index}" - ) - merged[key] = value + available = float(value) != 0.0 and all( + key in chunk and float(chunk[key]) != 0.0 for chunk in chunks[1:] + ) + merged[key] = np.float32(1.0 if available else 0.0) elif key == "sid" or value is None: merged[key] = value + elif any(key not in chunk for chunk in chunks[1:]): + continue elif key == "fid": merged[key] = [frame_id for chunk in chunks for frame_id in chunk[key]] else: merged[key] = np.concatenate([chunk[key] for chunk in chunks], axis=0) + for chunk in chunks[1:]: + for key in chunk: + if key.startswith("find_"): + merged.setdefault(key, np.float32(0.0)) return merged @@ -1359,6 +1843,38 @@ def _scan_frame_nlocs( return nlocs +def _group_positions_by_value(values: np.ndarray) -> dict[int, np.ndarray]: + """Group the positions of ``values`` by the value each holds. + + Sorting once and cutting at the value boundaries keeps the work inside + NumPy. Accumulating Python lists instead costs one interpreter iteration + and one boxed integer per element, which at the 10^8 frames a large LMDB + holds is minutes of runtime and tens of gigabytes of resident memory. + + Parameters + ---------- + values : numpy.ndarray + Integer values to group by, with shape ``(n,)``. + + Returns + ------- + dict[int, numpy.ndarray] + Value → the ascending positions holding it. Each array is a slice of + one shared buffer, so the mapping costs one extra array in total. + """ + if values.size == 0: + return {} + order = np.argsort(values, kind="stable") + ordered = values[order] + cuts = np.flatnonzero(ordered[1:] != ordered[:-1]) + 1 + starts = np.concatenate(([0], cuts)) + stops = np.concatenate((cuts, [values.size])) + return { + int(ordered[start]): order[start:stop] + for start, stop in zip(starts.tolist(), stops.tolist(), strict=True) + } + + def _compute_batch_size(nloc: int, rule: int) -> int: """Compute batch_size for a given nloc using the auto rule.""" bsi = rule // max(nloc, 1) @@ -1456,11 +1972,12 @@ def __init__( ) -> None: self.lmdb_path = str(Path(lmdb_path).resolve()) self._type_map = type_map + # Read before opening the frame-serving environment, which disables + # the readahead this one large sequential value depends on. Training + # draws frames in shuffled order and so leaves it disabled. + meta = _read_metadata_of(self.lmdb_path) self._env = _open_lmdb(self.lmdb_path) - with self._env.begin() as txn: - meta = _read_metadata(txn) - self.nframes, self._frame_fmt, self._natoms_per_type = _parse_metadata(meta) self._natoms = sum(self._natoms_per_type) self._ntypes = len(type_map) @@ -1499,13 +2016,23 @@ def __init__( # the dataset-to-original mapping lives in ``self._retained_keys``. # Metadata carries the counts when the writer recorded them; otherwise # each frame's atom_types shape is scanned (~10 us/frame). + # Both tables are held as NumPy arrays rather than Python lists: at + # 10^8 frames a list of boxed integers costs tens of gigabytes, while + # the arrays cost 4 bytes an entry and let the grouping below run as + # array operations. The list msgpack unpacks is released immediately. meta_nlocs = meta.get("frame_nlocs") if meta_nlocs is not None: - orig_frame_nlocs = [int(n) for n in meta_nlocs] + orig_frame_nlocs = np.fromiter( + meta_nlocs, dtype=np.int32, count=len(meta_nlocs) + ) else: - orig_frame_nlocs = _scan_frame_nlocs( - self._env, self.nframes, self._frame_fmt, self._natoms + orig_frame_nlocs = np.asarray( + _scan_frame_nlocs( + self._env, self.nframes, self._frame_fmt, self._natoms + ), + dtype=np.int32, ) + del meta_nlocs # Parse frame_system_ids for auto_prob support. ``_nsystems`` must stay # at ``max(original_sid) + 1`` even after filter:N so that user-facing @@ -1513,11 +2040,14 @@ def __init__( # keeps its meaning across filter thresholds. meta_sys_ids = meta.get("frame_system_ids") if meta_sys_ids is not None: - orig_frame_system_ids: list[int] | None = [int(s) for s in meta_sys_ids] - self._nsystems = max(orig_frame_system_ids) + 1 + orig_frame_system_ids: np.ndarray | None = np.fromiter( + meta_sys_ids, dtype=np.int32, count=len(meta_sys_ids) + ) + self._nsystems = int(orig_frame_system_ids.max()) + 1 else: orig_frame_system_ids = None self._nsystems = 1 + del meta_sys_ids, meta # Parse batch_size spec. ``auto_rule``, ``max_rule`` and ``mix_rule`` # are mutually exclusive; ``filter_rule`` implies ``max_rule`` plus @@ -1545,73 +2075,69 @@ def __init__( ) # Determine which original-index frames survive the filter. Without - # ``filter:N`` every frame is retained. + # ``filter:N`` every frame is retained, and the identity mapping is + # left as the ``arange`` so no frame is copied through a mask. + retained_keys = np.arange(self.nframes, dtype=np.int64) if self._filter_rule is not None: - retained_keys = [ - i for i, n in enumerate(orig_frame_nlocs) if n <= self._filter_rule - ] - n_dropped = self.nframes - len(retained_keys) + n_dropped = int(np.count_nonzero(orig_frame_nlocs > self._filter_rule)) if n_dropped > 0: + retained_keys = retained_keys[orig_frame_nlocs <= self._filter_rule] log.info( f"LMDB filter:{self._filter_rule} drops {n_dropped}/" f"{self.nframes} frames with nloc > {self._filter_rule} " f"({self.lmdb_path})." ) - else: - retained_keys = list(range(self.nframes)) # Dataset-index → original LMDB frame key. ``__getitem__`` looks up # this table so that ``reader[i]`` is a valid LMDB read for every # ``0 <= i < len(reader)``, no matter how many frames were filtered. - self._retained_keys: list[int] = retained_keys + self._retained_keys: np.ndarray = retained_keys # Re-key _frame_nlocs / _frame_system_ids into the dataset-index # space so that every downstream consumer (nloc_groups, system_groups, # LmdbBatchSampler, _expand_indices_by_blocks) operates in a single, # self-consistent indexing scheme. - self._frame_nlocs = [orig_frame_nlocs[k] for k in retained_keys] + keys_are_identity = retained_keys.size == self.nframes + self._keys_are_identity = keys_are_identity + self._frame_nlocs = ( + orig_frame_nlocs if keys_are_identity else orig_frame_nlocs[retained_keys] + ) - if orig_frame_system_ids is not None: - self._frame_system_ids: list[int] | None = [ - orig_frame_system_ids[k] for k in retained_keys - ] + if orig_frame_system_ids is None: + self._frame_system_ids: np.ndarray | None = None + elif keys_are_identity: + self._frame_system_ids = orig_frame_system_ids else: - self._frame_system_ids = None + self._frame_system_ids = orig_frame_system_ids[retained_keys] + + # nframes now reflects retained frames; __len__ returns this and the + # valid index domain for __getitem__ is [0, self.nframes). + self.nframes = int(retained_keys.size) # Group retained frames by nloc using dataset indices (0..len-1). # Statistics collection consumes these groups in every batching mode, # because per-nloc groups are the largest units that stack without # padding. - self._nloc_groups: dict[int, list[int]] = {} - for ds_idx, nloc in enumerate(self._frame_nlocs): - self._nloc_groups.setdefault(nloc, []).append(ds_idx) - - # Group retained frames by original system id; the sid numbering is - # preserved (no compression) so user-facing auto_prob slices stay - # meaningful across filter thresholds. Fully-dropped systems appear - # as zero-frame entries in ``_system_nframes``. + self._nloc_groups = _group_positions_by_value(self._frame_nlocs) + + # Frames per original system id; the sid numbering is preserved (no + # compression) so user-facing auto_prob slices stay meaningful across + # filter thresholds. Fully-dropped systems count zero. The per-system + # index lists themselves are built only if asked for; see + # :attr:`system_groups`. + self._system_groups: dict[int, np.ndarray] | None = None if self._frame_system_ids is not None: - self._system_groups: dict[int, list[int]] = {} - for ds_idx, sid in enumerate(self._frame_system_ids): - self._system_groups.setdefault(sid, []).append(ds_idx) - self._system_nframes: list[int] = [ - len(self._system_groups.get(i, [])) for i in range(self._nsystems) - ] + self._system_nframes: list[int] = np.bincount( + self._frame_system_ids, minlength=self._nsystems + ).tolist() else: - self._system_groups = {0: list(range(len(retained_keys)))} - self._system_nframes = [len(retained_keys)] - - # nframes now reflects retained frames; __len__ returns this and the - # valid index domain for __getitem__ is [0, self.nframes). - self.nframes = len(retained_keys) + self._system_nframes = [self.nframes] # Nominal batch size, reported to callers that want a single number. # The sampler never uses it: same-nloc modes go through # get_batch_size_for_nloc, and ``mix:N`` sizes each batch by budget. mean_nloc = ( - sum(self._frame_nlocs) / len(self._frame_nlocs) - if self._frame_nlocs - else self._natoms + float(self._frame_nlocs.mean()) if self._frame_nlocs.size else self._natoms ) if self._auto_rule is not None: self.batch_size = _compute_batch_size(self._natoms, self._auto_rule) @@ -1625,17 +2151,21 @@ def __init__( # Data requirements tracking self._data_requirements: dict[str, DataRequirementItem] = {} self._data_requirements_frozen = False + self._data_requirements_revision = 0 + # Requirements arrive after reader construction. Availability remains + # unresolved until the first grouping request so raw fields unused by + # the run cause no I/O and cannot influence the partition. + self._uniform_availability: bool | None = None + self._availability_index: _AvailabilityIndex | None = None self._decode_config = LmdbDecodeConfig( ntypes=self._ntypes, natoms=self._natoms, type_remap=self._type_remap, data_requirements=self._data_requirements, + dataset=self.lmdb_path, ) - # Availability signatures are decoded lazily and reused by every - # sampler epoch. Registering new requirements invalidates the cache. - self._find_signature_cache: dict[int, tuple[tuple[str, bool], ...]] = {} - # Which fields carry an atom axis follows from the requirements, so - # this cache is invalidated alongside the signature cache. + # Which fields carry an atom axis follows from the registered + # requirements, so this cache is invalidated when they change. self._per_atom_strides: dict[str, int] | None = None # Batches are rectangular until a consumer that reads a flat node axis # asks otherwise; see :meth:`use_ragged_batches`. @@ -1713,7 +2243,7 @@ def __getitem__(self, index: int) -> dict[str, Any]: self._data_requirements_frozen = True if index < 0 or index >= self.nframes: raise IndexError(f"dataset index {index} out of range [0, {self.nframes})") - original_key = self._retained_keys[index] + original_key = int(self._retained_keys[index]) key = format(original_key, self._frame_fmt).encode() raw = self._transaction().get(key) if raw is None: @@ -1736,7 +2266,7 @@ def original_keys(self, indices: Sequence[int]) -> list[int]: raise IndexError( f"dataset index {index} out of range [0, {self.nframes})" ) - keys.append(self._retained_keys[index]) + keys.append(int(self._retained_keys[index])) return keys def batch_pad_nloc(self, indices: Sequence[int]) -> int: @@ -1753,7 +2283,7 @@ def batch_pad_nloc(self, indices: Sequence[int]) -> int: The largest atom count in the batch. Batches drawn from a single nloc group return that group's atom count, so padding is a no-op. """ - return max(self._frame_nlocs[int(index)] for index in indices) + return int(self._frame_nlocs[np.asarray(indices, dtype=np.int64)].max()) def batch_layout( self, indices: Sequence[int], *, ragged: bool | None = None @@ -1782,10 +2312,11 @@ def batch_layout( The per-frame atom counts, the per-atom strides, and whether the frames are concatenated or padded to a common width. """ + # Widened after gathering, and only to the batch's own size: a layout + # carries int64 counts wherever it is built, including in the decode + # workers a chunk of it is shipped to. return BatchLayout.over( - np.asarray( - [self._frame_nlocs[int(index)] for index in indices], dtype=np.int64 - ), + self._frame_nlocs[np.asarray(indices, dtype=np.int64)].astype(np.int64), self.per_atom_strides(), ragged=self._ragged_batches if ragged is None else ragged, ) @@ -1890,74 +2421,146 @@ def closed(self) -> bool: # --- Data requirement interface --- + def _scan_exact_availability( + self, + availability_keys: Sequence[str], + ) -> _AvailabilityIndex: + """Build the exact index under sequential kernel readahead.""" + + def scan(transaction: lmdb.Transaction) -> _AvailabilityIndex: + def read_raw(position: int) -> bytes | None: + original_key = int(self._retained_keys[position]) + return transaction.get(format(original_key, self._frame_fmt).encode()) + + return _scan_availability_index( + self.nframes, + read_raw, + availability_keys, + self.lmdb_path, + ) + + environment = self._env + if environment is None: + raise RuntimeError("cannot scan a closed LMDB reader") + if bool(environment.flags().get("readahead", False)): + return scan(self._transaction()) + + resolved = str(Path(self.lmdb_path).resolve()) + cache_entry = _ENV_CACHE.get(resolved) + owns_environment_exclusively = ( + cache_entry is not None + and cache_entry[0] is environment + and cache_entry[1] == 1 + ) + if not owns_environment_exclusively: + if _is_local_rank_zero(): + log.info( + "LMDB label-availability scan uses an isolated sequential " + "reader because the random-read environment is shared: %s", + self.lmdb_path, + ) + source_index = _scan_lmdb_path_in_worker( + self.lmdb_path, + availability_keys, + ) + if self._keys_are_identity: + return source_index + return _AvailabilityIndex( + source_index.ids[self._retained_keys], + source_index.signature_count, + ) + + transaction = self._txn + if transaction is not None: + transaction.abort() + self._txn = None + self._env = None + _close_lmdb(self.lmdb_path) + try: + sequential_environment = _open_lmdb(self.lmdb_path, sequential=True) + with sequential_environment.begin() as sequential_transaction: + return scan(sequential_transaction) + finally: + _close_lmdb(self.lmdb_path) + self._env = _open_lmdb(self.lmdb_path, sequential=False) + self._txn = self._env.begin() + def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: - """Register expected keys; missing keys get default fill + find_key=0.0.""" + """Register the consumer data contract before resolving any frame.""" if self._data_requirements_frozen: raise RuntimeError( "LMDB data requirements must be registered before reading any frame" ) for item in data_requirement: self._data_requirements[item["key"]] = item - self._find_signature_cache.clear() + self._data_requirements_revision += 1 + self._uniform_availability = None + self._availability_index = None self._per_atom_strides = None - def get_find_signature(self, index: int) -> tuple[tuple[str, bool], ...]: - """Return the scalar availability signature for one retained frame. + def availability_groups(self, indices: np.ndarray) -> list[np.ndarray]: + """Partition dataset indices into label-compatible groups. - The signature covers registered data requirements and optional model - inputs whose ``find_*`` flags are created by :meth:`__getitem__`. - Reading only msgpack keys avoids decoding large coordinate and label - arrays while the sampler partitions frames. - """ - cached = self._find_signature_cache.get(index) - if cached is not None: - return cached - if index < 0 or index >= self.nframes: - raise IndexError(f"dataset index {index} out of range [0, {self.nframes})") + A ``find_*`` flag is one scalar per batch. Optional tracked labels + therefore remain homogeneous so present labels are not discarded. + Mandatory labels fail during frame decoding; default-backed and + derived fields do not participate. - original_key = self._retained_keys[index] - key = format(original_key, self._frame_fmt).encode() - raw = self._txn.get(key) - if raw is None: - raise IndexError( - f"Frame {original_key} not found in LMDB (dataset index {index})" - ) - raw_frame = msgpack.unpackb(raw, raw=False) - frame = {_KEY_REMAP.get(name, name): value for name, value in raw_frame.items()} - signature_keys = _availability_signature_keys( - frame, iter(self._data_requirements) - ) - signature = [] - for data_key in signature_keys: - find_key = f"find_{data_key}" - if find_key in frame: - find_value = _decode_value(frame[find_key]) - available = bool(float(np.asarray(find_value).item())) - elif data_key == "min_pair_dist" and data_key in self._data_requirements: - # __getitem__ computes this requirement when it is not stored. - available = True - else: - available = data_key in frame - signature.append((find_key, available)) + The bounded uniformity probe is deferred until requirements have been + registered. A dataset that probes mixed builds one compact exact index + for the entire retained dataset. Statistics and every sampler epoch + reuse that index without further LMDB reads. - result = tuple(signature) - self._find_signature_cache[index] = result - return result + The probe reads a bounded sample, so "uniform" is a finding and not a + proof; :data:`_AVAILABILITY_PROBE_FRAMES` states what a missed frame + costs and how a dataset can settle the question exactly. - def group_indices_by_find_signature( - self, indices: list[int] - ) -> dict[tuple[tuple[str, bool], ...], list[int]]: - """Partition dataset indices into scalar-compatible label groups.""" - groups: dict[tuple[tuple[str, bool], ...], list[int]] = {} - for index in indices: - groups.setdefault(self.get_find_signature(index), []).append(index) - return groups + Parameters + ---------- + indices : numpy.ndarray + Dataset indices to partition. + + Returns + ------- + list[numpy.ndarray] + Non-empty index groups in a stable order, together covering + ``indices``. + """ + if len(indices) == 0: + return [] + availability_keys = _availability_requirement_keys(self._data_requirements) + if not availability_keys: + self._uniform_availability = True + return [indices] + if self._uniform_availability is None: + self._uniform_availability = _probe_uniform_availability( + self._transaction(), + _evenly_spaced(self._retained_keys, _AVAILABILITY_PROBE_FRAMES), + self._frame_fmt, + availability_keys, + ) + if self._uniform_availability: + return [indices] + if self._availability_index is None: + self._availability_index = self._scan_exact_availability( + availability_keys, + ) + index_array = np.asarray(indices) + return self._availability_index.groups( + index_array, + positions=index_array, + ) @property def data_requirements(self) -> list[DataRequirementItem]: """Registered data requirements in insertion order.""" return list(self._data_requirements.values()) + @property + def data_requirements_revision(self) -> int: + """Monotonic revision used to invalidate dependent sampling plans.""" + return self._data_requirements_revision + def print_summary(self, name: str, prob: Any) -> None: """Print basic dataset info.""" n_groups = len(self._nloc_groups) @@ -2043,13 +2646,13 @@ def type_map(self) -> list[str]: return self._type_map @property - def nloc_groups(self) -> dict[int, list[int]]: - """Nloc → list of frame indices.""" + def nloc_groups(self) -> dict[int, np.ndarray]: + """Atom count → the ascending dataset indices holding it.""" return self._nloc_groups @property - def frame_nlocs(self) -> list[int]: - """Per-frame atom count.""" + def frame_nlocs(self) -> np.ndarray: + """Per-frame atom count, indexed by dataset index.""" return self._frame_nlocs @property @@ -2058,13 +2661,24 @@ def nsystems(self) -> int: return self._nsystems @property - def frame_system_ids(self) -> list[int] | None: - """Per-frame system index, or None if not available.""" + def frame_system_ids(self) -> np.ndarray | None: + """Per-frame system index, or None when the metadata omits it.""" return self._frame_system_ids @property - def system_groups(self) -> dict[int, list[int]]: - """System index → list of frame indices.""" + def system_groups(self) -> dict[int, np.ndarray]: + """System index → the ascending dataset indices of its frames. + + Built on demand: the per-system index lists cost an array the size of + the dataset, which the training and statistics paths never need -- + they consult :attr:`system_nframes` instead. + """ + if self._system_groups is None: + self._system_groups = ( + _group_positions_by_value(self._frame_system_ids) + if self._frame_system_ids is not None + else {0: np.arange(self.nframes, dtype=np.int64)} + ) return self._system_groups @property @@ -2073,6 +2687,44 @@ def system_nframes(self) -> list[int]: return self._system_nframes +def _batch_find_flags(frames: list[dict[str, Any]]) -> dict[str, np.float32]: + """Reduce the per-frame ``find_*`` flags to the one scalar a batch carries. + + A ``find_*`` flag is a single scalar for the whole batch, so a label + counts as available only where every frame supplies it. A frame that + lacks the label carries a default fill in its place, and reporting the + label as unavailable is what keeps that fill out of the loss. + + Registered data requirements are default-filled on every frame by + :func:`decode_lmdb_frame`. Availability-sensitive labels are normally + partitioned before collation to preserve their usable frames; this + reduction remains the correctness boundary when the bounded probe misses + a rare frame. Default-backed and derived requirements may mix + deliberately, while mandatory requirements fail during frame resolution. + + Parameters + ---------- + frames : list[dict[str, Any]] + Per-frame dicts about to be stacked. + + Returns + ------- + dict[str, numpy.float32] + One scalar flag for every ``find_*`` key that any frame carries. + """ + find_keys = sorted( + {key for frame in frames for key in frame if key.startswith("find_")} + ) + return { + key: np.float32( + 1.0 + if all(key in frame and float(frame[key]) != 0.0 for frame in frames) + else 0.0 + ) + for key in find_keys + } + + def _pad_atom_axis(xp: Any, array: Any, length: int, fill: int, device: Any) -> Any: """Widen one per-atom array's leading axis to ``length``.""" if array.shape[0] == length: @@ -2096,10 +2748,16 @@ def collate_lmdb_frames( etc. The array library is inferred from the first frame's ``coord``. Conventions match :func:`deepmd.dpmodel.utils.batch.normalize_batch`: - ``find_*`` flags remain scalar and must be constant within a batch; - ``fid`` is collected as a list; ``type`` is dropped (callers should - already use ``atype``); other arrays are stacked along axis 0. A ``sid`` - placeholder is appended. + ``find_*`` flags remain scalar; ``fid`` is collected as a list; ``type`` + is dropped (callers should already use ``atype``); other arrays are + stacked along axis 0. A ``sid`` placeholder is appended. + + Frames need not agree on label availability. A label only some of them + carry is reported unavailable for the whole batch, as + :func:`_batch_find_flags` describes, and a field only some of them carry + is left out of the batch entirely because it cannot be stacked. Neither + can happen to a registered data requirement, which + :func:`decode_lmdb_frame` puts on every frame. The batch keeps the key order of its frames, which is the order :func:`decode_lmdb_batch` also produces, so a batch is the same mapping @@ -2108,8 +2766,7 @@ def collate_lmdb_frames( Parameters ---------- frames : list[dict[str, Any]] - Per-frame dicts to stack, all sharing one label-availability - signature. + Per-frame dicts to stack. per_atom_keys : frozenset[str], optional Fields whose leading axis is the atom axis. When the frames differ in atom count these are padded to the batch maximum, with ``atype`` @@ -2130,23 +2787,7 @@ def collate_lmdb_frames( xp = array_api_compat.array_namespace(frames[0]["coord"]) dev = array_api_compat.device(frames[0]["coord"]) - # Availability must agree across the batch before the flags can collapse - # to one scalar per key. Frames are checked ahead of collation so a mixed - # batch is reported rather than silently reduced to its first frame. - find_keys = sorted( - {key for frame in frames for key in frame if key.startswith("find_")} - ) - for key in find_keys: - if any(key not in frame for frame in frames): - raise ValueError( - f"LMDB batch has inconsistent availability metadata for {key!r}" - ) - values = [float(frame[key]) for frame in frames] - if any(value != values[0] for value in values[1:]): - raise ValueError( - f"LMDB batch mixes {key!r} values {values}; " - "LmdbBatchSampler must group frames by label availability" - ) + find_flags = _batch_find_flags(frames) strides = per_atom_strides(frames[0], per_atom_keys) if per_atom_keys else {} pad_nloc = max(frame["coord"].shape[0] for frame in frames) if strides else 0 @@ -2154,13 +2795,18 @@ def collate_lmdb_frames( out: dict[str, Any] = {} for key in frames[0]: if key.startswith("find_"): - out[key] = frames[0][key] + out[key] = find_flags[key] elif key == "fid": out[key] = [f[key] for f in frames] elif key == "type": continue elif frames[0][key] is None: out[key] = None + elif any(key not in frame for frame in frames): + # A field only some frames carry cannot be stacked. Its ``find_`` + # flag is false by the same token, so the batch stays coherent + # without it. + continue elif key in strides: length = pad_nloc * strides[key] fill = _pad_fill_value(key) @@ -2169,6 +2815,10 @@ def collate_lmdb_frames( ) else: out[key] = xp.stack([f[key] for f in frames]) + # A flag a later frame raised that the first one never had still belongs + # to the batch, reporting the label as unavailable. + for key, flag in find_flags.items(): + out.setdefault(key, flag) out["sid"] = xp.asarray([0], dtype=xp.int64, device=dev) return out @@ -2311,7 +2961,7 @@ def _expand_indices_by_blocks( indices : list[int] Frame indices in the current nloc group. frame_system_ids : np.ndarray - Per-frame system id for the entire dataset (int64 array). + Per-frame system id for the entire dataset. block_targets : list[tuple[list[int], int]] Per-block (system_ids, total_target_frames). rng : np.random.Generator @@ -2320,12 +2970,12 @@ def _expand_indices_by_blocks( Pre-computed total actual frame count per block (across all nloc groups). When provided, avoids an O(N) scan of frame_system_ids. _sid_to_blk_arr : np.ndarray or None - Pre-computed system-id to block-index lookup array. When provided, + Pre-computed lookup from :func:`system_block_lookup`. When provided, avoids rebuilding the mapping for each call. _group_block_targets : list[int] or None Exact target for each block in this group. Production samplers - allocate these targets globally across all ``(nloc, find-signature)`` - groups so independent rounding cannot change a block's total size. + allocate these targets globally across all groups so independent + rounding cannot change a block's total size. Returns ------- @@ -2334,30 +2984,21 @@ def _expand_indices_by_blocks( """ n_blocks = len(block_targets) - # Build sys_id -> block_idx lookup array if _sid_to_blk_arr is None: - sys_to_block: dict[int, int] = {} - for blk_idx, (sys_ids, _target) in enumerate(block_targets): - for sid in sys_ids: - sys_to_block[sid] = blk_idx - max_sid = max(sys_to_block.keys()) + 1 if sys_to_block else 0 - _sid_to_blk_arr = np.full(max_sid, -1, dtype=np.int32) - for sid, blk in sys_to_block.items(): - _sid_to_blk_arr[sid] = blk - - # Partition indices by block using numpy for speed + _sid_to_blk_arr = system_block_lookup(block_targets) + + sid_arr = np.asarray(frame_system_ids) idx_arr = np.asarray(indices, dtype=np.int64) - sid_arr = np.asarray(frame_system_ids, dtype=np.int64) - # Vectorized lookup: get block id for each index - idx_sids = sid_arr[idx_arr] - idx_blks = _sid_to_blk_arr[idx_sids] + idx_blks = resolve_frame_blocks(idx_arr, sid_arr, _sid_to_blk_arr) # Pre-compute block_total_actual if not provided if _block_total_actual is None and _group_block_targets is None: - _block_total_actual = [] - for sys_ids, _ in block_targets: - total = sum(int(np.sum(sid_arr == sid)) for sid in sys_ids) - _block_total_actual.append(total) + _block_total_actual = count_group_blocks( + np.arange(sid_arr.size, dtype=np.int64), + sid_arr, + _sid_to_blk_arr, + n_blocks, + ).tolist() expanded_parts: list[np.ndarray] = [] @@ -2417,9 +3058,13 @@ def _expand_indices_by_blocks( def collect_lmdb_sampling_groups( reader: "LmdbDataReader", -) -> list[tuple[int, list[int]]]: +) -> list[tuple[int, np.ndarray]]: """Collect homogeneous LMDB groups shared by training and statistics. + Atom count comes from metadata and costs no read. Availability-sensitive + requirements subdivide a group only when the bounded probe detects mixed + source presence; see :meth:`LmdbDataReader.availability_groups`. + Parameters ---------- reader : LmdbDataReader @@ -2427,27 +3072,22 @@ def collect_lmdb_sampling_groups( Returns ------- - list[tuple[int, list[int]]] + list[tuple[int, numpy.ndarray]] Stable ``(nloc, frame indices)`` groups compatible with collation. """ - groups: list[tuple[int, list[int]]] = [] - for nloc in sorted(reader.nloc_groups): - signature_groups = reader.group_indices_by_find_signature( - list(reader.nloc_groups[nloc]) - ) - for signature in sorted(signature_groups): - groups.append((nloc, list(signature_groups[signature]))) - return groups + return [ + (nloc, group) + for nloc in sorted(reader.nloc_groups) + for group in reader.availability_groups(reader.nloc_groups[nloc]) + ] -def _collect_batch_groups(reader: "LmdbDataReader") -> list[list[int]]: +def _collect_batch_groups(reader: "LmdbDataReader") -> list[np.ndarray]: """Collect the groups a training batch may be drawn from. - A group is the largest set of frames one batch may span. Label - availability always partitions it, because ``find_*`` flags collapse to a - single scalar per batch. Atom count partitions it as well in every mode - but ``mix:N``, whose decoded batch accommodates unequal counts and so - needs only the availability split. + A group is the largest set of frames one batch may span. Atom count + partitions the frames in every mode but ``mix:N``, whose decoded batch + accommodates unequal counts and so is bounded only by label availability. Parameters ---------- @@ -2456,21 +3096,108 @@ def _collect_batch_groups(reader: "LmdbDataReader") -> list[list[int]]: Returns ------- - list[list[int]] + list[numpy.ndarray] Frame indices per group, in the stable order shared by iteration and length. """ if reader.mixed_nloc: - signature_groups = reader.group_indices_by_find_signature( - list(range(len(reader))) - ) - return [list(signature_groups[key]) for key in sorted(signature_groups)] + return reader.availability_groups(np.arange(len(reader), dtype=np.int64)) return [indices for _nloc, indices in collect_lmdb_sampling_groups(reader)] +def system_block_lookup( + block_targets: list[tuple[list[int], int]], +) -> np.ndarray: + """Build the system-id to block-index lookup a whole group indexes at once. + + The table carries one row beyond the highest system id named by a block, + holding the "no block" marker. Clamping a system id into the table then + maps everything a block does not name onto that row, which keeps the + lookup total without a per-frame membership test. + + Parameters + ---------- + block_targets : list[tuple[list[int], int]] + Per-block ``(system_ids, target_frame_count)``. + + Returns + ------- + numpy.ndarray + Block index of each system id, ``-1`` where no block names it, with + shape ``(max_system_id + 2,)``. + """ + max_sid = max( + (sid for system_ids, _target in block_targets for sid in system_ids), + default=-1, + ) + lookup = np.full(max_sid + 2, -1, dtype=np.int64) + for block_index, (system_ids, _target) in enumerate(block_targets): + if system_ids: + lookup[np.asarray(system_ids, dtype=np.int64)] = block_index + return lookup + + +def resolve_frame_blocks( + indices: np.ndarray | list[int], + frame_system_ids: np.ndarray, + lookup: np.ndarray, +) -> np.ndarray: + """Map each frame of one group to the block it belongs to. + + Parameters + ---------- + indices : numpy.ndarray or list[int] + Dataset indices of one group. + frame_system_ids : numpy.ndarray + Per-frame system id of the whole dataset. Gathered from, never + widened: it holds one entry per frame, so converting its dtype would + copy the entire dataset for the sake of a single group. + lookup : numpy.ndarray + System-id to block-index table from :func:`system_block_lookup`. + + Returns + ------- + numpy.ndarray + Block index of each frame, ``-1`` where no block claims it. + """ + sids = np.asarray(frame_system_ids)[np.asarray(indices, dtype=np.int64)] + return lookup[np.minimum(sids, lookup.size - 1)] + + +def count_group_blocks( + indices: np.ndarray | list[int], + frame_system_ids: np.ndarray, + lookup: np.ndarray, + n_blocks: int, +) -> np.ndarray: + """Count how many of one group's frames fall in each block. + + Parameters + ---------- + indices : numpy.ndarray or list[int] + Dataset indices of one group. + frame_system_ids : numpy.ndarray + Per-frame system id of the whole dataset. + lookup : numpy.ndarray + System-id to block-index table from :func:`system_block_lookup`. + n_blocks : int + Number of blocks. + + Returns + ------- + numpy.ndarray + Frame count per block, with shape ``(n_blocks,)``. Frames belonging + to no block are excluded, so the counts need not sum to the group + size. + """ + blocks = resolve_frame_blocks(indices, frame_system_ids, lookup) + # Shift by one so that the "no block" marker lands in bin zero. + return np.bincount(blocks + 1, minlength=n_blocks + 1)[1:] + + def _allocate_group_block_targets( - groups: list[list[int]], - frame_system_ids: list[int] | np.ndarray, + groups: list[np.ndarray], + frame_system_ids: np.ndarray, block_targets: list[tuple[list[int], int]], ) -> list[list[int]]: """Allocate every block target exactly across homogeneous groups. @@ -2480,17 +3207,13 @@ def _allocate_group_block_targets( remainder method. Stable group order breaks equal-remainder ties, which keeps distributed ranks deterministic without floating-point rounding. """ - group_actual = [[0] * len(block_targets) for _ in groups] - system_to_block = { - system_id: block_index - for block_index, (system_ids, _target) in enumerate(block_targets) - for system_id in system_ids - } - for group_index, indices in enumerate(groups): - for index in indices: - block_index = system_to_block.get(int(frame_system_ids[index])) - if block_index is not None: - group_actual[group_index][block_index] += 1 + lookup = system_block_lookup(block_targets) + group_actual = [ + count_group_blocks( + indices, frame_system_ids, lookup, len(block_targets) + ).tolist() + for indices in groups + ] group_targets = [counts.copy() for counts in group_actual] for block_index, (_system_ids, block_target) in enumerate(block_targets): @@ -2529,17 +3252,20 @@ def _allocate_group_block_targets( return group_targets -def _chop_same_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[int]]: +def _chop_same_nloc( + reader: "LmdbDataReader", indices: Sequence[int] +) -> list[list[int]]: """Split one homogeneous group into fixed-size batches.""" - batch_size = reader.get_batch_size_for_nloc(reader.frame_nlocs[indices[0]]) + batch_size = reader.get_batch_size_for_nloc(int(reader.frame_nlocs[indices[0]])) + index_list = np.asarray(indices).tolist() return [ - indices[start : start + batch_size] - for start in range(0, len(indices), batch_size) + index_list[start : start + batch_size] + for start in range(0, len(index_list), batch_size) ] def _chop_mixed_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[int]]: - """Split one availability group into batches under an atom-axis budget. + """Split one group into batches under an atom-axis budget. ``mix:N`` budgets the length of a batch's atom axis, and the layout the batch will be decoded in decides both how that length is measured and the @@ -2590,7 +3316,7 @@ def _chop_mixed_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[ reader : LmdbDataReader Provides the per-frame atom counts, the atom budget and the layout. indices : list[int] - Dataset indices of one label-availability group. + Dataset indices of one group. Returns ------- @@ -2602,7 +3328,9 @@ def _chop_mixed_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[ raise ValueError("mixed-nloc batching requires a batch_size of 'mix:N'") index_array = np.asarray(indices, dtype=np.int64) - nloc_array = np.asarray(reader.frame_nlocs, dtype=np.int64)[index_array] + # Gathered, not widened: the atom counts are one entry per frame, so a + # dtype conversion here would copy the whole dataset for one group. + nloc_array = reader.frame_nlocs[index_array] if not reader.ragged_batches: order = np.argsort(nloc_array, kind="stable") index_array, nloc_array = index_array[order], nloc_array[order] @@ -2630,7 +3358,7 @@ def _build_all_batches( rng: np.random.Generator, block_targets: list[tuple[list[int], int]] | None = None, ) -> list[list[int]]: - """Build batches homogeneous in label availability. + """Build the batches of one pass over the dataset. Groups are chopped into batches, then interleaved round-robin so that consecutive batches come from different groups. Under ``mix:N`` a batch @@ -2652,7 +3380,7 @@ def _build_all_batches( Returns ------- list[list[int]] - Each inner list has one scalar ``find_*`` signature. + Dataset indices grouped into batches. """ groups = _collect_batch_groups(reader) chop = _chop_mixed_nloc if reader.mixed_nloc else _chop_same_nloc @@ -2665,25 +3393,16 @@ def _build_all_batches( sid_to_blk_arr: np.ndarray | None = None group_block_targets: list[list[int]] | None = None if block_targets and reader.frame_system_ids is not None: - # Convert frame_system_ids to numpy once - sid_arr = np.array(reader.frame_system_ids, dtype=np.int64) + sid_arr = reader.frame_system_ids group_block_targets = _allocate_group_block_targets( groups, sid_arr, block_targets ) - # Build sys_id -> block_idx lookup array once - sys_to_block: dict[int, int] = {} - for blk_idx, (sys_ids, _target) in enumerate(block_targets): - for sid in sys_ids: - sys_to_block[sid] = blk_idx - max_sid = max(sys_to_block.keys()) + 1 if sys_to_block else 0 - sid_to_blk_arr = np.full(max_sid, -1, dtype=np.int32) - for sid, blk in sys_to_block.items(): - sid_to_blk_arr[sid] = blk + sid_to_blk_arr = system_block_lookup(block_targets) for group_index, original_indices in enumerate(groups): indices = original_indices - # Expand each availability group independently using targets that - # were allocated globally, preserving both scalar flags and totals. + # Expand each group independently using targets that were allocated + # globally, so that a block's total size is preserved exactly. if block_targets and sid_arr is not None and group_block_targets is not None: indices = _expand_indices_by_blocks( indices, @@ -2694,8 +3413,10 @@ def _build_all_batches( _group_block_targets=group_block_targets[group_index], ) if shuffle: - rng.shuffle(indices) - group_batches.append(chop(reader, indices) if indices else []) + # ``permutation`` returns a copy, which matters because a group + # may alias one of the reader's own index tables. + indices = rng.permutation(indices) + group_batches.append(chop(reader, indices) if len(indices) else []) # Interleave groups round-robin all_batches: list[list[int]] = [] @@ -2713,14 +3434,17 @@ def _build_all_batches( class LmdbBatchSampler: - """Batch sampler over an LMDB, grouped by ``find_*`` signature. + """Batch sampler over an LMDB, grouped by atom count. + + Atom count is handled by the reader's batching rule: all rules but + ``mix:N`` keep a batch uniform in atom count, while ``mix:N`` fills + batches to the atom budget its decoded layout is measured against. + Groups are interleaved round-robin and the batch order is then shuffled, + so training sees a varied mix. - Every batch carries one label-availability signature, because ``find_*`` - flags collapse to a single scalar per batch. Atom count is handled by the - reader's batching rule: all rules but ``mix:N`` additionally keep a batch - uniform in atom count, while ``mix:N`` fills batches to the atom budget - its decoded layout is measured against. Groups are interleaved round-robin - and the batch order is then shuffled, so training sees a varied mix. + Optional tracked labels subdivide these groups when a bounded probe + detects mixed presence. Default-backed and derived fields may mix freely; + mandatory fields fail when an unavailable frame is decoded. The sampler serves one pass at a time, drawn from ``seed + epoch``. The pending pass is materialized before it is served, which is what lets @@ -2756,6 +3480,7 @@ def __init__( self._epoch = 0 self._block_targets = block_targets self._batches: list[list[int]] | None = None + self._data_requirements_revision = -1 def batches(self) -> list[list[int]]: """Return the batch list of the pending pass, building it if needed. @@ -2763,10 +3488,10 @@ def batches(self) -> list[list[int]]: Returns ------- list[list[int]] - Dataset indices grouped into batches, one scalar ``find_*`` - signature each. + Dataset indices grouped into batches. """ - if self._batches is None: + revision = self._reader.data_requirements_revision + if self._batches is None or self._data_requirements_revision != revision: seed = None if self._seed is None else self._seed + self._epoch self._batches = _build_all_batches( self._reader, @@ -2774,6 +3499,7 @@ def batches(self) -> list[list[int]]: np.random.default_rng(seed), self._block_targets, ) + self._data_requirements_revision = revision return self._batches def set_epoch(self, epoch: int) -> None: @@ -2788,15 +3514,6 @@ def set_epoch(self, epoch: int) -> None: self._epoch = epoch self._batches = None - def refresh_batch_count(self) -> None: - """Discard the pending pass after the frame grouping changed. - - The pass is materialized ahead of iteration so that ``__len__`` can - report it exactly, which leaves it stale once new data requirements - repartition the frames by label availability. - """ - self._batches = None - def __iter__(self) -> Iterator[list[int]]: """Yield the pending pass, and move the epoch on to its successor.""" batches = self.batches() @@ -2865,10 +3582,6 @@ def __init__( self._block_targets = block_targets self._global: LmdbBatchSampler | None = None - def refresh_batch_count(self) -> None: - """Discard the pending global pass after the frame grouping changed.""" - self._global = None - def set_epoch(self, epoch: int) -> None: """Set epoch for deterministic cross-rank shuffling. @@ -3022,13 +3735,16 @@ def __init__( ) -> None: self.lmdb_path = str(lmdb_path) self._type_map = type_map or [] - self._env = _open_lmdb(self.lmdb_path) - - with self._env.begin() as txn: - meta = _read_metadata(txn) + meta = _read_metadata_of(self.lmdb_path) + # Every read this class serves walks a group in ascending key order, + # which is the pattern kernel readahead exists for. + self._env = _open_lmdb(self.lmdb_path, sequential=True) self.nframes, self._frame_fmt, self._natoms_per_type = _parse_metadata(meta) self._natoms = sum(self._natoms_per_type) + self._ntypes = ( + len(self._type_map) if self._type_map else len(self._natoms_per_type) + ) # Build type remapping if LMDB's type_map differs from model's type_map lmdb_type_map = meta.get("type_map") @@ -3060,6 +3776,18 @@ def __init__( # Data requirements self._requirements: dict[str, dict[str, Any]] = {} + self._uniform_availability: bool | None = None + self._availability_indices: dict[ + int, tuple[np.ndarray, _AvailabilityIndex] + ] = {} + self._full_availability_index: _AvailabilityIndex | None = None + self._decode_config = LmdbDecodeConfig( + ntypes=self._ntypes, + natoms=self._natoms, + type_remap=self._type_remap, + data_requirements=self._requirements, + dataset=self.lmdb_path, + ) # Detect PBC from the first retained frame. self.pbc = True @@ -3079,7 +3807,7 @@ def _select_frames( meta: dict, shuffle_test: bool, max_frames: float | None, - ) -> dict[int, list[int]]: + ) -> dict[int, np.ndarray]: """Group the frame indices by atom count, then sample each group. Parameters @@ -3095,50 +3823,43 @@ def _select_frames( Returns ------- - dict[int, list[int]] + dict[int, numpy.ndarray] The retained LMDB frame indices of each atom count. """ raw_nlocs = meta.get("frame_nlocs") if _is_encoded_array(raw_nlocs): - nlocs = _decode_array(raw_nlocs).reshape(-1).astype(np.int64) + nlocs = _decode_array(raw_nlocs).reshape(-1).astype(np.int32) elif raw_nlocs is not None: - nlocs = np.asarray(raw_nlocs, dtype=np.int64) + nlocs = np.fromiter(raw_nlocs, dtype=np.int32, count=len(raw_nlocs)) else: nlocs = np.asarray( _scan_frame_nlocs( self._env, self.nframes, self._frame_fmt, self._natoms ), - dtype=np.int64, + dtype=np.int32, ) - # Sorting once groups every atom count in a single pass, which matters - # for datasets whose frame count reaches into the millions. - order = np.argsort(nlocs, kind="stable") - starts = np.concatenate( - ([0], np.flatnonzero(np.diff(nlocs[order])) + 1, [nlocs.size]) - ) - + groups = _group_positions_by_value(nlocs) keep = ( None if max_frames is None or not np.isfinite(max_frames) else int(max_frames) ) - groups: dict[int, list[int]] = {} - for begin, end in pairwise(starts): - indices = order[begin:end].copy() + if not shuffle_test and keep is None: + return groups + + for nloc, indices in groups.items(): if shuffle_test: dp_random.shuffle(indices) - if keep is not None: - indices = indices[:keep] - groups[int(nlocs[order[begin]])] = indices.tolist() + groups[nloc] = indices if keep is None else indices[:keep] return groups - def _read_frames(self, frame_indices: list[int]) -> list[dict[str, Any]]: + def _read_frames(self, frame_indices: Sequence[int]) -> list[dict[str, Any]]: """Decode the given LMDB frames, applying the type remapping. Parameters ---------- - frame_indices : list[int] + frame_indices : Sequence[int] Indices of the frames to read, as keyed in the LMDB. Returns @@ -3147,77 +3868,121 @@ def _read_frames(self, frame_indices: list[int]) -> list[dict[str, Any]]: One decoded frame per index that the LMDB holds. """ frames: list[dict[str, Any]] = [] - with self._env.begin() as txn: - for index in frame_indices: - raw = txn.get(format(index, self._frame_fmt).encode()) + with self._env.begin() as transaction: + for index in np.asarray(frame_indices).tolist(): + raw = transaction.get(format(index, self._frame_fmt).encode()) if raw is None: continue - frame = _remap_keys(_decode_frame(raw)) - atype = frame.get("atype") - if self._type_remap is not None and isinstance(atype, np.ndarray): - frame["atype"] = _remap_atom_types( - atype.reshape(-1), self._type_remap + frames.append( + decode_lmdb_frame( + raw, + int(index), + self._decode_config, + copy_arrays=True, ) - frames.append(frame) + ) return frames def __del__(self) -> None: - """Release the LMDB environment ref-count on garbage collection.""" - path = getattr(self, "lmdb_path", None) - if path is not None: - _close_lmdb(path) + """Release the LMDB environment ref-count on garbage collection. + + The count is released only once, and only if construction got as far + as taking it: an instance that failed earlier holds no reference, and + releasing one it never took would close the environment underneath + whichever reader does hold it. + """ + if getattr(self, "_env", None) is None: + return + self._env = None + _close_lmdb(self.lmdb_path) @property - def nloc_groups(self) -> dict[int, list[int]]: + def nloc_groups(self) -> dict[int, np.ndarray]: """Nloc → the LMDB frame indices retained for that atom count.""" return self._nloc_groups - @staticmethod - def _frame_has_data(frame: dict[str, Any], key: str) -> bool: - """Resolve one frame's explicit or inferred ``find_*`` value. + def availability_groups(self, frame_indices: np.ndarray) -> list[np.ndarray]: + """Partition LMDB frame indices into label-compatible groups. - The frame may still carry its msgpack payload, so that availability - can be settled without decoding the arrays it describes. - """ - find_key = f"find_{key}" - if find_key in frame: - return bool(float(np.asarray(_decode_value(frame[find_key])).item())) - value = frame.get(key) - if _is_encoded_array(value): - return True - return isinstance(value, (np.ndarray, np.generic, int, float, bool)) + The validation counterpart of + :meth:`LmdbDataReader.availability_groups`. Only optional tracked + labels participate. Exact signature IDs are cached per retained nloc + group and reused by every full-validation pass. - @property - def find_signature_groups( - self, - ) -> dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]]: - """Group the retained frame indices by atom count and label availability. + Parameters + ---------- + frame_indices : numpy.ndarray + LMDB frame indices to partition. - Frames that :meth:`_stack_frames` would refuse to stack together land - in different groups, because both settle availability with - :meth:`_frame_has_data`. Only the msgpack payload of each frame is - read, so the grouping does not decode the arrays it separates. + Returns + ------- + list[numpy.ndarray] + Non-empty index groups in a stable order, together covering the + input. An empty group would name no frames to stack, so an empty + input yields no group at all. """ - groups: dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]] = {} - with self._env.begin() as transaction: - for nloc, frame_indices in self._nloc_groups.items(): - for index in frame_indices: - raw = transaction.get(format(index, self._frame_fmt).encode()) - if raw is None: - continue - frame = _remap_keys(msgpack.unpackb(raw, raw=False)) - signature = tuple( - (f"find_{key}", self._frame_has_data(frame, key)) - for key in _availability_signature_keys( - frame, iter(self._requirements) - ) + if not len(frame_indices): + return [] + availability_keys = _availability_requirement_keys(self._requirements) + if not availability_keys: + self._uniform_availability = True + return [frame_indices] + if self._uniform_availability is None: + groups = list(self._nloc_groups.values()) + per_group = max(1, _AVAILABILITY_PROBE_FRAMES // max(len(groups), 1)) + with self._env.begin() as transaction: + self._uniform_availability = _probe_uniform_availability( + transaction, + ( + key + for group in groups + for key in _evenly_spaced(group, per_group) + ), + self._frame_fmt, + availability_keys, + ) + if self._uniform_availability: + return [frame_indices] + index_array = np.asarray(frame_indices) + if not bool(self._env.flags().get("readahead", False)): + if self._full_availability_index is None: + self._full_availability_index = _scan_lmdb_path_in_worker( + self.lmdb_path, + availability_keys, + ) + return self._full_availability_index.groups( + index_array, + positions=index_array, + ) + + cache_key = id(frame_indices) + cached = self._availability_indices.get(cache_key) + if cached is None or cached[0] is not frame_indices: + with self._env.begin() as transaction: + + def read_raw(position: int) -> bytes | None: + frame_index = int(index_array[position]) + return transaction.get( + format(frame_index, self._frame_fmt).encode() ) - groups.setdefault((nloc, signature), []).append(index) - return groups - def get_test_by_indices(self, frame_indices: list[int]) -> dict[str, Any]: + availability_index = _scan_availability_index( + len(index_array), + read_raw, + availability_keys, + self.lmdb_path, + ) + self._availability_indices[cache_key] = ( + frame_indices, + availability_index, + ) + else: + availability_index = cached[1] + return availability_index.groups(index_array) + + def get_test_by_indices(self, frame_indices: Sequence[int]) -> dict[str, Any]: """Stack one homogeneous validation group selected by frame index.""" - if not frame_indices: + if not len(frame_indices): raise ValueError("frame_indices must contain at least one frame") frames = self._read_frames(frame_indices) nlocs = { @@ -3242,18 +4007,25 @@ def add( repeat: int = 1, default: float = 0.0, dtype: np.dtype | None = None, + source_policy: DataRequirementSourcePolicy = "tracked", **kwargs: Any, ) -> None: """Register a data requirement (mirrors DeepmdData.add).""" - self._requirements[key] = { - "ndof": ndof, - "atomic": atomic, - "must": must, - "high_prec": high_prec, - "repeat": repeat, - "default": default, - "dtype": dtype, - } + requirement = DataRequirementItem( + key, + ndof, + atomic=atomic, + must=must, + high_prec=high_prec, + repeat=repeat, + default=default, + dtype=dtype, + source_policy=source_policy, + ) + self._requirements[key] = requirement.dict + self._uniform_availability = None + self._availability_indices.clear() + self._full_availability_index = None def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: """Register expected keys from ``DataRequirementItem`` objects. @@ -3272,6 +4044,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N repeat=item["repeat"], default=item["default"], dtype=item["dtype"], + source_policy=item["source_policy"], ) def _resolve_dtype(self, key: str) -> np.dtype: @@ -3308,6 +4081,7 @@ def iter_test( chunk_atoms: int, numb_test: float = float("inf"), nloc: int | None = None, + frame_indices: Sequence[int] | None = None, ) -> Iterator[dict[str, Any]]: """Yield the test frames in chunks, reading each chunk on demand. @@ -3324,13 +4098,18 @@ def iter_test( serves every frame of the group. nloc : int or None, optional Atom count selecting the group, resolved as in :meth:`get_test`. + frame_indices : Sequence[int] or None, optional + Frames to serve in place of the whole group, which is how a + label-compatible subgroup names itself. They must all have the + atom count ``nloc`` selects. Yields ------ dict[str, Any] One chunk of frames, stacked as :meth:`get_test` stacks them. """ - frame_indices, natoms = self._resolve_group(nloc) + group_indices, natoms = self._resolve_group(nloc) + frame_indices = group_indices if frame_indices is None else frame_indices if np.isfinite(numb_test): frame_indices = frame_indices[: int(numb_test)] step = max(1, int(chunk_atoms) // max(1, natoms)) @@ -3338,7 +4117,7 @@ def iter_test( chunk = frame_indices[begin : begin + step] yield self._stack_frames(self._read_frames(chunk), natoms) - def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]: + def _resolve_group(self, nloc: int | None) -> tuple[np.ndarray, int]: """Return the retained frame indices and atom count of one group. Parameters @@ -3349,7 +4128,7 @@ def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]: Returns ------- - tuple[list[int], int] + tuple[numpy.ndarray, int] The LMDB frame indices of the group and its atom count. Raises @@ -3378,7 +4157,12 @@ def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]: def _stack_frames( self, frames: list[dict[str, Any]], natoms: int ) -> dict[str, Any]: - """Stack a list of same-nloc frames into numpy arrays.""" + """Stack a list of same-nloc frames into numpy arrays. + + Frames need not agree on label availability; a label only some of + them carry is reported unavailable for the whole group, mirroring + :func:`_batch_find_flags` on the training path. + """ nframes = len(frames) result: dict[str, Any] = {} @@ -3418,7 +4202,9 @@ def _stack_frames( # Dynamically discover all data keys from the first frame, plus # any registered requirements. Structural keys (coord, box, type) # are excluded — they are already handled above. - _structural_keys = frozenset({"coord", "box", "atype"}) + _structural_keys = frozenset( + {"coord", "box", "atype", "natoms", "real_natoms_vec", "fid"} + ) all_keys: dict[str, dict[str, Any]] = {} if frames: for fk in frames[0]: @@ -3429,61 +4215,29 @@ def _stack_frames( for key, req in self._requirements.items(): all_keys[key] = req - for key, req_info in all_keys.items(): - availability = [self._frame_has_data(frame, key) for frame in frames] - if any(flag != availability[0] for flag in availability[1:]): - raise ValueError( - f"LMDB validation group mixes find_{key} values {availability}" - ) - has_key = availability[0] - if not has_key and req_info.get("must", False): - raise RuntimeError(f"Required LMDB test-data field {key!r} is missing.") - result[f"find_{key}"] = 1.0 if has_key else 0.0 - - # Get repeat factor from registered requirements - repeat = 1 - if key in self._requirements: - repeat = self._requirements[key].get("repeat", 1) - - if has_key: - arrays = [] - for frame in frames: - val = frame.get(key) - if isinstance(val, np.ndarray): - arr = val.astype(self._resolve_dtype(key)).ravel() - if repeat != 1: - arr = np.repeat(arr, repeat) - arrays.append(arr) - elif val is not None: - arrays.append( - np.array([float(val)], dtype=self._resolve_dtype(key)) - ) - else: - ref = next( - ( - f[key] - for f in frames - if isinstance(f.get(key), np.ndarray) - ), - None, - ) - if ref is not None: - size = ref.size * repeat if repeat != 1 else ref.size - arrays.append( - np.zeros(size, dtype=self._resolve_dtype(key)) - ) - else: - arrays.append(np.zeros(1, dtype=self._resolve_dtype(key))) - result[key] = np.stack(arrays) - elif key in self._requirements: - ndof = self._requirements[key]["ndof"] - atomic = self._requirements[key]["atomic"] - default = self._requirements[key]["default"] - if atomic: - shape = (nframes, natoms * ndof * repeat) + for key in all_keys: + has_key = all(_frame_source_available(frame, key) for frame in frames) + requirement = self._requirements.get(key) + if requirement is None and not has_key: + # An unregistered field that only part of the group carries + # has no complete column and remains outside the result. + continue + arrays: list[np.ndarray] = [] + for frame in frames: + value = frame.get(key) + if isinstance(value, np.ndarray): + arrays.append(value.astype(self._resolve_dtype(key)).ravel()) + elif value is not None: + arrays.append( + np.array([float(value)], dtype=self._resolve_dtype(key)) + ) else: - shape = (nframes, ndof * repeat) - result[key] = np.full(shape, default, dtype=self._resolve_dtype(key)) + raise RuntimeError( + f"Resolved LMDB field {key!r} is absent in frame " + f"{frame.get('fid', '')} of {self.lmdb_path}" + ) + result[key] = np.stack(arrays) + result[f"find_{key}"] = 1.0 if has_key else 0.0 return result @@ -3492,18 +4246,18 @@ class LmdbTestDataNlocView: """Expose one stack-compatible subset of :class:`LmdbTestData`. The underlying :class:`LmdbTestData` groups frames by atom count. This - view fixes one ``nloc`` and can additionally select a homogeneous - label-availability subgroup. All other attributes (``pbc``, - ``mixed_type``, …) are forwarded to the underlying object. It lets - downstream consumers that expect a ``DeepmdData``-style system work on - mixed-nloc or partially labeled LMDB datasets without vector find flags. + view fixes one ``nloc`` and can additionally select a label-compatible + subgroup within it. All other attributes (``pbc``, ``mixed_type``, …) + are forwarded to the underlying object. It lets downstream consumers + that expect a ``DeepmdData``-style system work on mixed-nloc or + partially labeled LMDB datasets without vector find flags. """ def __init__( self, lmdb_test_data: "LmdbTestData", nloc: int, - frame_indices: list[int] | None = None, + frame_indices: Sequence[int] | None = None, ) -> None: self._inner = lmdb_test_data self._nloc = nloc @@ -3523,18 +4277,12 @@ def iter_test( chunk_atoms: int, numb_test: float = float("inf"), ) -> Iterator[dict[str, Any]]: - """Yield this group's frames in chunks.""" - if self._frame_indices is not None: - frame_indices = self._frame_indices - if np.isfinite(numb_test): - frame_indices = frame_indices[: int(numb_test)] - step = max(1, int(chunk_atoms) // max(1, self._nloc)) - return ( - self._inner.get_test_by_indices(frame_indices[begin : begin + step]) - for begin in range(0, len(frame_indices), step) - ) + """Yield this group's frames in chunks, as :meth:`get_test` selects them.""" return self._inner.iter_test( - chunk_atoms=chunk_atoms, numb_test=numb_test, nloc=self._nloc + chunk_atoms=chunk_atoms, + numb_test=numb_test, + nloc=self._nloc, + frame_indices=self._frame_indices, ) diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 381d07d97c..7e10e99466 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -535,7 +535,8 @@ def forward( find_atom_pref = ( label.get("find_atom_pref", 0.0) if not self.use_default_pf else 1.0 ) - pref_pf = pref_pf * find_atom_pref + effective_find_pf = find_force * find_atom_pref + pref_pf = pref_pf * effective_find_pf atom_pref_reshape = atom_pref.reshape(-1) if self.loss_func == "mse": @@ -544,7 +545,7 @@ def forward( ).mean() if not self.inference: more_loss["l2_pref_force_loss"] = self.display_if_exist( - l2_pref_force_loss.detach(), find_atom_pref + l2_pref_force_loss.detach(), effective_find_pf ) if maskf is not None: # Idiom 1 with pref weight (ncomp=3). @@ -556,7 +557,7 @@ def forward( loss += (pref_pf * l2_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) rmse_pf = l2_pf_masked.sqrt() more_loss["rmse_pf"] = self.display_if_exist( - rmse_pf.detach(), find_atom_pref + rmse_pf.detach(), effective_find_pf ) else: loss += (pref_pf * l2_pref_force_loss).to( @@ -564,7 +565,7 @@ def forward( ) rmse_pf = l2_pref_force_loss.sqrt() more_loss["rmse_pf"] = self.display_if_exist( - rmse_pf.detach(), find_atom_pref + rmse_pf.detach(), effective_find_pf ) elif self.loss_func == "mae": l1_pref_force_loss = (torch.abs(diff_f) * atom_pref_reshape).mean() @@ -576,14 +577,14 @@ def forward( ) loss += (pref_pf * l1_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) more_loss["mae_pf"] = self.display_if_exist( - l1_pf_masked.detach(), find_atom_pref + l1_pf_masked.detach(), effective_find_pf ) else: loss += (pref_pf * l1_pref_force_loss).to( GLOBAL_PT_FLOAT_PRECISION ) more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss.detach(), find_atom_pref + l1_pref_force_loss.detach(), effective_find_pf ) else: raise NotImplementedError( @@ -593,7 +594,8 @@ def forward( if self.has_gf and "drdq" in label: drdq = label["drdq"] find_drdq = label.get("find_drdq", 0.0) - pref_gf = pref_gf * find_drdq + effective_find_gf = find_force * find_drdq + pref_gf = pref_gf * effective_find_gf if maskf is not None: # Mask per-atom forces before projecting onto generalized coords. f_3d = force_pred.reshape(_nf, _nloc, 3) * maskf.reshape( @@ -627,12 +629,12 @@ def forward( l2_gen_force_loss = torch.square(diff_gen_force).mean() if not self.inference: more_loss["l2_gen_force_loss"] = self.display_if_exist( - l2_gen_force_loss.detach(), find_drdq + l2_gen_force_loss.detach(), effective_find_gf ) loss += (pref_gf * l2_gen_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) rmse_gf = l2_gen_force_loss.sqrt() more_loss["rmse_gf"] = self.display_if_exist( - rmse_gf.detach(), find_drdq + rmse_gf.detach(), effective_find_gf ) if self.has_v and "virial" in model_pred and "virial" in label: @@ -891,6 +893,7 @@ def label_requirement(self) -> list[DataRequirementItem]: high_prec=False, repeat=3, default=1.0, + source_policy="default" if self.use_default_pf else "tracked", ) ) if self.has_gf > 0: @@ -912,6 +915,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, default=1.0, + source_policy="default", ) ) if self.has_h: diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index fa324315e7..701cf62228 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -211,6 +211,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, default=1.0, + source_policy="default", ) ) return label_requirement diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 1101127e1a..876a174e14 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -507,6 +507,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: must=False, high_prec=False, default=min_pair_dist, + source_policy="derived", ) ) training_data.add_data_requirement(data_requirement) @@ -585,6 +586,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: must=False, high_prec=False, default=min_pair_dist, + source_policy="derived", ) ) training_data[model_key].add_data_requirement(data_requirement) @@ -2354,6 +2356,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=False, must=not _model.has_default_fparam(), default=_fparam_default, + source_policy=("default" if _model.has_default_fparam() else "tracked"), ) ] additional_data_requirement += fparam_requirement_items @@ -2382,6 +2385,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=True, must=not allow_missing_spin, default=0.0, + source_policy="default" if allow_missing_spin else "tracked", ) ] additional_data_requirement += spin_requirement_items @@ -2397,6 +2401,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=False, must=not has_default_cs, default=cs_default, + source_policy="default" if has_default_cs else "tracked", ) ) return additional_data_requirement diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 888d08692f..1648adcd04 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -24,9 +24,12 @@ LmdbDecodeConfig, LmdbTestData, collate_lmdb_frames, + collect_lmdb_sampling_groups, compute_block_targets, + count_group_blocks, is_lmdb, resolve_per_atom_keys, + system_block_lookup, ) from deepmd.env import ( get_lmdb_num_workers, @@ -60,8 +63,9 @@ def _collate_lmdb_batch( via ``array_api_compat``). Frames of different atom counts are padded to the batch maximum; the - padded slots carry the phantom atom type. Frames must still agree on - label availability, which :class:`LmdbBatchSampler` guarantees. + padded slots carry the phantom atom type. Frames need not agree on label + availability: a label only some of them carry is reported unavailable + for the whole batch. Parameters ---------- @@ -248,37 +252,31 @@ def __init__( collate_fn=self._collate, ) - # Per-nloc and label-availability dataloaders for make_stat_input. - # These are rebuilt after requirements are registered so statistics - # never collate real labels with default-filled placeholders. + # Homogeneous dataloaders for make_stat_input, built on first use and + # discarded whenever new requirements change how frames decode. self._nloc_dataloaders: list[DataLoader] | None = None - def _rebuild_nloc_dataloaders(self) -> None: - """Build homogeneous loaders used by model-stat collection.""" + def _build_nloc_dataloaders(self) -> None: + """Build the homogeneous loaders used by model-stat collection.""" dataloaders: list[DataLoader] = [] - for nloc in sorted(self._reader.nloc_groups.keys()): - signature_groups = self._reader.group_indices_by_find_signature( - self._reader.nloc_groups[nloc] - ) - for signature in sorted(signature_groups): - subset = torch.utils.data.Subset(self, signature_groups[signature]) - bs = self._reader.get_batch_size_for_nloc(nloc) - with torch.device("cpu"): - dl = DataLoader( - subset, - batch_size=bs, - shuffle=False, - num_workers=0, - drop_last=False, - collate_fn=self._collate, - ) - dataloaders.append(dl) + for nloc, indices in collect_lmdb_sampling_groups(self._reader): + subset = torch.utils.data.Subset(self, indices) + with torch.device("cpu"): + dl = DataLoader( + subset, + batch_size=self._reader.get_batch_size_for_nloc(nloc), + shuffle=False, + num_workers=0, + drop_last=False, + collate_fn=self._collate, + ) + dataloaders.append(dl) self._nloc_dataloaders = dataloaders def _get_nloc_dataloaders(self) -> list[DataLoader]: """Materialize statistics loaders lazily when none are registered.""" if self._nloc_dataloaders is None: - self._rebuild_nloc_dataloaders() + self._build_nloc_dataloaders() dataloaders = self._nloc_dataloaders if dataloaders is None: raise RuntimeError("Failed to initialize LMDB statistics dataloaders") @@ -324,7 +322,10 @@ def data_requirements(self) -> list[DataRequirementItem]: def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: self._reader.add_data_requirement(data_requirement) - self._rebuild_nloc_dataloaders() + # Loaders decode through the registered requirements, so any already + # built are stale. They are rebuilt on demand, which spares a run + # whose statistics come from a stat file the work entirely. + self._nloc_dataloaders = None def close(self) -> None: """Release parent-process LMDB resources.""" @@ -365,11 +366,12 @@ def print_summary(self, name: str, prob: Any) -> None: f"{actual}->{target} (x{ratio:.2f})" ) - # Build sys_id -> block_idx mapping - sys_to_block: dict[int, int] = {} - for blk_idx, (sys_ids, _) in enumerate(self._block_targets): - for sid in sys_ids: - sys_to_block[sid] = blk_idx + # A whole nloc group's block membership resolves as one indexing + # operation, which a dataset of 10^8 frames needs it to be. The + # lookup is the one the sampler allocates its targets with, so + # both agree on which systems a block claims. + n_blocks = len(self._block_targets) + lookup = system_block_lookup(self._block_targets) # Compute expanded nloc counts analytically (no actual expansion) expanded_nloc_info = {} @@ -377,19 +379,12 @@ def print_summary(self, name: str, prob: Any) -> None: if reader.frame_system_ids is None: expanded_nloc_info[nloc] = len(indices) continue - # Count indices per block in this nloc group - blk_counts: dict[int, int] = {} - unassigned = 0 - for idx in indices: - sid = reader.frame_system_ids[idx] - blk = sys_to_block.get(sid) - if blk is not None: - blk_counts[blk] = blk_counts.get(blk, 0) + 1 - else: - unassigned += 1 - expanded = unassigned + counts = count_group_blocks( + indices, reader.frame_system_ids, lookup, n_blocks + ) + expanded = len(indices) - int(counts.sum()) for blk_idx, (_, blk_target) in enumerate(self._block_targets): - n_actual = blk_counts.get(blk_idx, 0) + n_actual = int(counts[blk_idx]) if n_actual == 0: continue bta = block_total_actual[blk_idx] @@ -439,8 +434,8 @@ def systems(self) -> list: def dataloaders(self) -> list: """Homogeneous dataloaders for make_stat_input. - Each loader has one nloc and one availability signature, so stat - collection sees consistent shapes and scalar ``find_*`` flags. + Each loader draws from one atom count and one label availability, so + stat collection sees consistent shapes and scalar ``find_*`` flags. """ return self._get_nloc_dataloaders() diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index fb9a2393e0..e1769a117d 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -346,6 +346,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=False, must=not has_default_fparam, default=fparam_default, + source_policy="default" if has_default_fparam else "tracked", ) ) if _model.get_dim_aparam() > 0: @@ -368,6 +369,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=True, must=not allow_missing_spin, default=0.0, + source_policy="default" if allow_missing_spin else "tracked", ) ) if _model.has_chg_spin_ebd(): @@ -387,6 +389,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: atomic=False, must=not has_default_cs, default=default_cs, + source_policy="default" if has_default_cs else "tracked", ) ) return additional_data_requirement @@ -1698,6 +1701,11 @@ def __init__( else self.valid_numb_batch_by_task[DEFAULT_TASK_KEY] ) + # The layout settles before the schedule below, because under + # ``mix:N`` it is what decides how many frames a batch holds, and the + # schedule counts those batches to turn epochs into steps. + self._configure_batch_layout(training_data, validation_data) + # Statistics ---------------------------------------------------------- self._finetune_update_stat = False self._sample_funcs: dict[str, Any] = {} @@ -2093,7 +2101,6 @@ def update_finetune_bias( self._configure_neighbor_graph_method( training_params.get("neighbor_graph_method", "auto") ) - self._configure_batch_layout(training_data, validation_data) # torch.compile ------------------------------------------------------- if self.enable_compile: @@ -2210,6 +2217,12 @@ def _configure_batch_layout(self, *data_maps: Any) -> None: each task's own: a multi-task run pairing a graph model with a dense one gives the first concatenated batches and the second padded ones. + This must run before the run length is resolved. Under ``mix:N`` the + layout decides where the sampler cuts a batch -- padding is what makes + a batch's cost depend on its widest frame -- so a schedule that + counted batches first would count a packing training never uses, and + would serve its first epoch in that packing. + Parameters ---------- *data_maps : Any diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index edd89e6d04..554294abca 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -432,10 +432,10 @@ def _iter_validation_data_systems(self) -> Iterator[Any]: - An LMDB-backed dataset owns an ``_reader``. Its frames are lazily materialized into a :class:`LmdbTestData` snapshot (cached across calls) and yielded as one :class:`LmdbTestDataNlocView` per - atom-count and label-availability group. Grouping by atom count lets - mixed-nloc frames be stacked, and grouping by label availability - keeps the scalar ``find_*`` flags valid so default-filled labels stay - out of the metrics. + atom-count and label-availability group. Grouping by atom count + lets mixed-nloc frames be stacked, and grouping by availability + keeps the scalar ``find_*`` flags valid so default-filled labels + stay out of the metrics. - A ``DeepmdDataSystem`` owns ``data_systems``, which are already ``DeepmdData`` instances. - A loader set owns ``systems``, each wrapping a ``DeepmdData`` in @@ -444,10 +444,11 @@ def _iter_validation_data_systems(self) -> Iterator[Any]: validation_data = self.validation_data if hasattr(validation_data, "_reader"): lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data) - for (nloc, _signature), indices in sorted( - lmdb_test_data.find_signature_groups.items() - ): - yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices) + for nloc in sorted(lmdb_test_data.nloc_groups): + for indices in lmdb_test_data.availability_groups( + lmdb_test_data.nloc_groups[nloc] + ): + yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices) return if hasattr(validation_data, "data_systems"): diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index d645ff9955..9af91be82c 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -119,7 +119,7 @@ def __init__( ) def _refresh_stat_groups(self) -> None: - """Rebuild statistical systems from the training sampler's groups.""" + """Build statistical systems from stack-compatible reader groups.""" self._stat_groups = collect_lmdb_sampling_groups(self._reader) self._stat_offsets = [0] * len(self._stat_groups) @@ -156,7 +156,7 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: Parameters ---------- sys_idx : int - Index into the ``(nloc, label-availability)`` groups. + Index into the per-atom-count groups. Returns ------- @@ -208,12 +208,8 @@ def get_stat_numb_batches(self, sys_idx: int) -> int: def add_data_requirements( self, data_requirement: list[DataRequirementItem] ) -> None: - # Batches are partitioned by label availability, so new requirements - # repartition the frames. Both the statistical groups and the pass the - # sampler holds pending are therefore rebuilt from the new partition. self._reader.add_data_requirement(data_requirement) self._refresh_stat_groups() - self._sampler.refresh_batch_count() def close(self) -> None: """Cancel prefetched work and release decoder processes.""" diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 4bde3a8f9e..a232aab393 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -18,6 +18,7 @@ ) from typing import ( Any, + Literal, ) import numpy as np @@ -35,6 +36,8 @@ log = logging.getLogger(__name__) +DataRequirementSourcePolicy = Literal["tracked", "default", "derived"] + class DeepmdData: """Class for a data system. @@ -1198,6 +1201,13 @@ class DataRequirementItem: special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + source_policy : {"tracked", "default", "derived"}, optional + How source availability affects the consumer. ``"tracked"`` keeps + source presence in the ``find_*`` contract. ``"default"`` treats the + configured default as valid input and reports the resolved field + available. ``"derived"`` computes the value from structural frame + data. Only optional tracked fields require availability-homogeneous + batching. """ def __init__( @@ -1213,7 +1223,18 @@ def __init__( dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, special_shape: str | None = None, + source_policy: DataRequirementSourcePolicy = "tracked", ) -> None: + if source_policy not in {"tracked", "default", "derived"}: + raise ValueError( + "source_policy must be 'tracked', 'default', or 'derived', " + f"got {source_policy!r}" + ) + if must and source_policy != "tracked": + raise ValueError( + f"mandatory data requirement {key!r} cannot use " + f"source_policy={source_policy!r}" + ) self.key = key self.ndof = ndof self.atomic = atomic @@ -1225,6 +1246,7 @@ def __init__( self.dtype = dtype self.output_natoms_for_type_sel = output_natoms_for_type_sel self.special_shape = special_shape + self.source_policy = source_policy self.dict = self.to_dict() def to_dict(self) -> dict: @@ -1239,6 +1261,7 @@ def to_dict(self) -> dict: "default": self.default, "dtype": self.dtype, "output_natoms_for_type_sel": self.output_natoms_for_type_sel, + "source_policy": self.source_policy, } if self.special_shape is not None: data["special_shape"] = self.special_shape diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 3537c97d15..ad8d5cf670 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -4,6 +4,7 @@ Pure dpmodel (NumPy/lmdb) tests — no PyTorch dependency. """ +import gc import os import signal import subprocess @@ -44,6 +45,8 @@ _expand_indices_by_blocks, _merge_lmdb_chunks, _remap_atom_types, + collate_lmdb_frames, + collect_lmdb_sampling_groups, compute_block_targets, decode_lmdb_batch, decode_lmdb_frame, @@ -462,6 +465,70 @@ def test_requirements_are_rejected_after_the_first_decode(self): reader.add_data_requirement(requirement) reader.close() + def test_mandatory_fields_reject_missing_and_explicitly_unavailable_sources(self): + """Mandatory LMDB inputs fail at the shared frame-resolution boundary.""" + path = _create_lmdb( + f"{self._tmpdir.name}/mandatory_spin.lmdb", + nframes=3, + natoms=6, + ) + environment = lmdb.open(path, readonly=False, lock=False) + with environment.begin(write=True) as transaction: + unavailable_key = format(1, "012d").encode() + unavailable_frame = msgpack.unpackb( + transaction.get(unavailable_key), raw=False + ) + unavailable_frame["spin"] = { + "type": " None: + """The declared labels must match the terms evaluated by the loss.""" + loss_fn = EnergyLoss(starter_learning_rate=1.0) + required_keys = {item.key for item in loss_fn.label_requirement} + self.assertEqual(required_keys, {"energy", "force"}) + + model_dict, all_labels, natoms = self._make_data() + label_dict = { + key: value + for key, value in all_labels.items() + if key in required_keys + or (key.startswith("find_") and key.removeprefix("find_") in required_keys) + } + loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict) + + self.assertTrue(np.isfinite(loss)) + self.assertEqual(set(more_loss), {"rmse_e", "rmse_f", "rmse"}) + + def test_prefactor_force_declares_both_inputs(self) -> None: + """Atomic force weighting requires force and atom_pref labels.""" + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + ) + self.assertEqual( + {item.key for item in loss_fn.label_requirement}, + {"force", "atom_pref"}, + ) + + def test_default_prefactor_is_valid_training_data(self) -> None: + """use_default_pf makes a missing atom_pref equivalent to unit weights.""" + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + use_default_pf=True, + ) + atom_pref = next( + item for item in loss_fn.label_requirement if item.key == "atom_pref" + ) + self.assertEqual(atom_pref.source_policy, "default") + def test_forward(self) -> None: loss_fn = EnergyLoss( starter_learning_rate=1.0, @@ -141,6 +188,56 @@ def test_forward(self) -> None: self.assertIn("rmse_gf", more_loss) self.assertIn("rmse_pf", more_loss) + def test_force_derived_terms_require_force_labels(self) -> None: + """PF and GF supervision are unavailable when force is unavailable.""" + cases = ( + ( + "prefactor force", + { + "start_pref_pf": 1.0, + "limit_pref_pf": 1.0, + "use_default_pf": True, + }, + "rmse_pf", + 0, + ), + ( + "generalized force", + { + "start_pref_gf": 1.0, + "limit_pref_gf": 1.0, + "numb_generalized_coord": 2, + }, + "rmse_gf", + 2, + ), + ) + for name, term_kwargs, metric, numb_generalized_coord in cases: + with self.subTest(term=name): + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + **term_kwargs, + ) + model_dict, label_dict, natoms = self._make_data( + numb_generalized_coord=numb_generalized_coord, + ) + label_dict["find_force"] = 0.0 + loss, more_loss = loss_fn.call( + 1.0, + natoms, + model_dict, + label_dict, + ) + + self.assertEqual(float(loss), 0.0) + self.assertTrue(np.isnan(more_loss[metric])) + class TestEnergyLossHuber(TestEnergyLossBase): """Test energy loss with Huber loss.""" diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index 2010b57016..d8f70d0dd3 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -35,6 +35,9 @@ LmdbDataset, _collate_lmdb_batch, ) +from deepmd.pt.utils.stat import ( + make_stat_input, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -551,9 +554,10 @@ def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatc sampler = LmdbBatchSampler(ds._reader, shuffle=True, seed=11) batches = list(sampler) assert len(sampler) == len(batches) == 2 + # Complementary labels must not share a batch: the frames carrying + # energy are indices 0 and 2, those carrying force are 1 and 3. for indices in batches: - signatures = {ds._reader.get_find_signature(index) for index in indices} - assert len(signatures) == 1 + assert len({index % 2 for index in indices}) == 1 distributed_batches = [] for rank in range(2): @@ -635,8 +639,8 @@ def zero_model(_batch=batch, **kwargs): assert force_loss.item() > 0.0 assert observed_flags == {(1.0, 0.0), (0.0, 1.0)} - def test_unrequested_labels_form_homogeneous_batches(self, tmp_path): - """Raw labels outside the loss requirements must still partition batches.""" + def test_unrequested_labels_do_not_partition_batches(self, tmp_path): + """Raw labels outside the active requirements cannot alter sampling.""" path = str(tmp_path / "partial-virial.lmdb") _create_partially_virial_lmdb(path) ds = LmdbDataset(path, type_map=["O", "H"], batch_size=2) @@ -644,10 +648,57 @@ def test_unrequested_labels_form_homogeneous_batches(self, tmp_path): [DataRequirementItem("energy", 1, atomic=False, must=False)] ) - with torch.device("cpu"): - batches = list(ds._inner_dataloader) - assert len(batches) == 2 - assert sorted("find_virial" in batch for batch in batches) == [False, True] + groups = ds._reader.availability_groups(np.arange(len(ds), dtype=np.int64)) + assert len(groups) == 1 + np.testing.assert_array_equal(groups[0], np.arange(len(ds))) + + def test_default_backed_fparam_survives_statistics(self, lmdb_dir): + """Statistics retain explicit and default-resolved frame parameters.""" + environment = lmdb.open(lmdb_dir, readonly=False, lock=False) + with environment.begin(write=True) as transaction: + for index in range(0, 10, 2): + key = format(index, "012d").encode() + frame = msgpack.unpackb(transaction.get(key), raw=False) + frame["fparam"] = { + "type": " bool: + """Return whether a scalar metric is NaN on any tensor device.""" + if isinstance(value, torch.Tensor): + return bool(torch.isnan(value.detach()).item()) + return bool(np.isnan(value)) + + class TestEnerStdLossDefaultPf(unittest.TestCase): """Test use_default_pf feature in EnergyStdLoss.""" @@ -176,6 +183,71 @@ def fake_model(): # The pref_force_loss should be a valid number (not NaN) self.assertFalse(np.isnan(pt_more_loss["l2_pref_force_loss"])) + def test_default_pf_requires_force_labels(self) -> None: + """Default atom weights cannot make a missing force label usable.""" + loss_fn = EnergyStdLoss( + self.start_lr, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + use_default_pf=True, + ) + label = dict(self.label_without_pref) + label["find_force"] = 0.0 + + _, loss, more_loss = loss_fn( + {}, + lambda: self.model_pred, + label, + self.nloc, + self.cur_lr, + ) + + self.assertEqual(float(loss.detach().cpu()), 0.0) + self.assertTrue(_is_nan(more_loss["l2_pref_force_loss"])) + self.assertTrue(_is_nan(more_loss["rmse_pf"])) + + def test_generalized_force_requires_force_labels(self) -> None: + """A present projection matrix cannot validate a missing force label.""" + numb_generalized_coord = 2 + loss_fn = EnergyStdLoss( + self.start_lr, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=numb_generalized_coord, + ) + label = dict(self.label_with_pref) + label["find_force"] = 0.0 + label["drdq"] = torch.ones( + (1, self.nloc * 3 * numb_generalized_coord), + dtype=label["force"].dtype, + device=label["force"].device, + ) + label["find_drdq"] = 1.0 + + _, loss, more_loss = loss_fn( + {}, + lambda: self.model_pred, + label, + self.nloc, + self.cur_lr, + ) + + self.assertEqual(float(loss.detach().cpu()), 0.0) + self.assertTrue(_is_nan(more_loss["l2_gen_force_loss"])) + self.assertTrue(_is_nan(more_loss["rmse_gf"])) + def test_default_pf_disabled(self) -> None: """With use_default_pf=False (default), pf loss should NOT be computed without find_atom_pref.""" loss_fn = EnergyStdLoss( @@ -287,6 +359,7 @@ def test_label_requirement_atom_pref_default(self) -> None: label_req = loss_fn.label_requirement atom_pref_req = next(r for r in label_req if r.key == "atom_pref") self.assertEqual(atom_pref_req.default, 1.0) + self.assertEqual(atom_pref_req.source_policy, "default") def test_serialize_deserialize(self) -> None: """Serialization round-trip should preserve use_default_pf.""" diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index e16a08172a..865cf027f3 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -963,6 +963,7 @@ def test_allow_missing_label_data_requirement( spin_req = next(rr for rr in reqs if rr.key == "spin") assert spin_req.must is expected_must assert spin_req.default == 0.0 + assert spin_req.source_policy == ("default" if allow_missing else "tracked") class TestPublicBaseModelRoundTrip: diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index fcc93e547b..ab48c8b339 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -21,6 +21,7 @@ import msgpack import numpy as np +from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module from deepmd.dpmodel.utils.batch import ( normalize_batch, split_batch, @@ -31,6 +32,12 @@ from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) +from deepmd.pt_expt.loss import ( + EnergyLoss, +) +from deepmd.pt_expt.train.training import ( + Trainer, +) from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) @@ -262,6 +269,27 @@ def test_data_requirements_freeze_after_first_read(self) -> None: ): ds.add_data_requirements([DataRequirementItem("late_label", ndof=1)]) + def test_missing_mandatory_model_input_is_rejected(self) -> None: + """A required model input cannot silently become a zero-filled batch.""" + ds = LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + num_workers=0, + ) + ds.add_data_requirements( + [DataRequirementItem("spin", 3, atomic=True, must=True)] + ) + try: + with self.assertRaisesRegex( + RuntimeError, + r"spin.*frame \d+.*test\.lmdb.*field is absent", + ): + ds.get_batch() + finally: + ds.close() + def test_get_batch_iterates_past_end(self) -> None: """get_batch reseeds the sampler at the end of an epoch.""" ds = LmdbDataSystem( @@ -327,6 +355,118 @@ def test_add_data_requirements_passthrough(self) -> None: self.assertIn("energy", batch) self.assertIn("find_energy", batch) + def test_default_atom_pref_does_not_partition_training_data(self) -> None: + """OC20-style unit defaults keep mixed atom_pref frames compatible.""" + path = os.path.join(self.tmpdir, "default_atom_pref.lmdb") + _create_test_lmdb(path, nframes=8, natoms=6) + environment = lmdb.open(path, readonly=False, lock=False) + with environment.begin(write=True) as transaction: + for index in range(0, 8, 2): + key = format(index, "012d").encode() + frame = msgpack.unpackb(transaction.get(key), raw=False) + frame["atom_pref"] = _encode_array(np.full(6, 2.0)) + transaction.put(key, msgpack.packb(frame, use_bin_type=True)) + environment.close() + + loss = EnergyLoss( + starter_learning_rate=0.002, + start_pref_e=20.0, + limit_pref_e=20.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_pf=20.0, + limit_pref_pf=20.0, + start_pref_v=5.0, + limit_pref_v=5.0, + loss_func="mae", + f_use_norm=True, + use_default_pf=True, + ) + with patch.object( + lmdb_data_module, + "_raw_frame_availability", + wraps=lmdb_data_module._raw_frame_availability, + ) as inspect_availability: + ds = LmdbDataSystem( + lmdb_path=path, + type_map=["O", "H"], + batch_size="mix:30000", + seed=0, + ) + self.assertEqual(inspect_availability.call_count, 0) + ds.add_data_requirements(loss.label_requirement) + + try: + self.assertLessEqual(inspect_availability.call_count, 8) + self.assertEqual(len(ds._stat_groups), 1) + batch = ds._reader.decode_batch([0, 1]) + np.testing.assert_array_equal(batch["atom_pref"][0], np.full(18, 2.0)) + np.testing.assert_array_equal(batch["atom_pref"][1], np.ones(18)) + self.assertEqual(float(batch["find_atom_pref"]), 1.0) + model_output = { + "energy": np.zeros_like(batch["energy"]), + "force": np.zeros_like(batch["force"]), + "virial": np.zeros_like(batch["virial"]), + "atom_energy": np.zeros((2, 6, 1), dtype=batch["energy"].dtype), + } + loss_value, metrics = loss.call(0.002, 6, model_output, batch) + self.assertTrue(np.isfinite(loss_value)) + self.assertIn("mae_pf", metrics) + finally: + ds.close() + + def test_missing_force_disables_default_prefactor_force_loss(self) -> None: + """Default atom weights cannot supervise a frame without force.""" + path = os.path.join(self.tmpdir, "partial_force_for_pf.lmdb") + _create_partially_labeled_lmdb(path) + loss = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + use_default_pf=True, + ) + ds = LmdbDataSystem( + lmdb_path=path, + type_map=["O", "H"], + batch_size=2, + seed=0, + num_workers=0, + ) + ds.add_data_requirements(loss.label_requirement) + try: + observed: dict[float, float] = {} + for _ in range(2): + batch = ds.get_batch() + nframes, nloc = batch["atype"].shape + model_output = { + "energy": np.zeros_like(batch.get("energy", np.zeros(nframes))), + "force": np.zeros_like(batch["force"]), + "virial": np.zeros((nframes, 9), dtype=batch["force"].dtype), + "atom_energy": np.zeros( + (nframes, nloc, 1), + dtype=batch["force"].dtype, + ), + } + loss_value, metrics = loss.call( + 1.0, + nloc, + model_output, + batch, + ) + find_force = float(batch["find_force"]) + observed[find_force] = float(loss_value) + if find_force == 0.0: + self.assertTrue(np.isnan(metrics["rmse_pf"])) + + self.assertEqual(observed[0.0], 0.0) + self.assertGreater(observed[1.0], 0.0) + finally: + ds.close() + def test_partial_labels_are_batched_by_availability(self) -> None: from deepmd.utils.data import ( DataRequirementItem, @@ -755,6 +895,43 @@ def test_dense_model_keeps_padded_batches(self) -> None: self.assertEqual(batch["atype"].ndim, 2) self.assertNotIn("n_node", batch) + def test_layout_is_settled_before_the_pass_is_counted(self) -> None: + """The run length must count the packing training actually uses. + + Under ``mix:N`` the layout decides where the sampler cuts a batch, and + the sampler materializes a pass the moment its length is asked for. + Settling the layout afterwards would count one packing and train on + another, and would serve the first epoch from the counted one. + + The pass is cached once built, so both orderings report the same count + after the fact. What distinguishes them is the layout in force at the + moment the count is taken, which is what this observes. + """ + config = self._config(self._dpa1()) + del config["training"]["numb_steps"] + config["training"]["numb_epoch"] = 1.0 + + observed: dict[str, bool] = {} + original = Trainer._epoch_length + + def record_layout(trainer_self, model_key: str) -> int: + observed["ragged_when_counted"] = ( + trainer_self.training_data._reader.ragged_batches + ) + return original(trainer_self, model_key) + + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + with patch.object(Trainer, "_epoch_length", record_layout): + trainer = get_trainer(config) + finally: + os.chdir(cwd) + + # DPA1 reads a flat node axis, so training runs on ragged batches. + self.assertTrue(trainer.training_data._reader.ragged_batches) + self.assertTrue(observed["ragged_when_counted"]) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 0569f21851..1c5e0aec3f 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -997,6 +997,7 @@ def has_chg_spin_ebd(self) -> bool: self.assertEqual(fparam_req.key, "fparam") self.assertEqual(fparam_req.ndof, 2) self.assertFalse(fparam_req.must) + self.assertEqual(fparam_req.source_policy, "default") # default is the model's default_fparam, not 0.0 self.assertNotIsInstance(fparam_req.default, float) import numpy as np @@ -1034,6 +1035,7 @@ def has_chg_spin_ebd(self) -> bool: self.assertEqual(fparam_req.key, "fparam") self.assertTrue(fparam_req.must) self.assertEqual(fparam_req.default, 0.0) + self.assertEqual(fparam_req.source_policy, "tracked") class TestRestart(unittest.TestCase): From ae387ea2a76aaf22a5cd06839d85b7c7025549db Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 9 Aug 2026 22:18:29 +0800 Subject: [PATCH 3/9] fix(lmdb): preserve seed and phantom edge filtering --- deepmd/pt/model/model/sezm_model.py | 22 ++++++++++++++++------ deepmd/pt/train/training.py | 12 +++++++----- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 0d907f1c6c..de2fb58d58 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -1529,14 +1529,18 @@ def core_compute( if comm_dict is None: descriptor_atype = atype + descriptor_real_atom = real_atom elif extended_atype is None: raise ValueError("`extended_atype` is required with `comm_dict`.") else: # Ghost atoms carry the type of the local atom they image, so this # sanitizes the phantoms among them on the same grounds as above. - descriptor_atype, _ = self._sanitize_atom_types(extended_atype) + descriptor_atype, descriptor_real_atom = self._sanitize_atom_types( + extended_atype + ) inter_potential_edge_mask = self._make_inter_potential_edge_mask( descriptor_atype, + descriptor_real_atom, edge_index, edge_mask, ) @@ -3241,6 +3245,7 @@ def reset_head_for_mode(self, mode: str) -> None: def _make_inter_potential_edge_mask( self, atype: torch.Tensor, + real_atom: torch.Tensor, edge_index: torch.Tensor, edge_mask: torch.Tensor, ) -> torch.Tensor: @@ -3249,7 +3254,9 @@ def _make_inter_potential_edge_mask( Parameters ---------- atype - Atom types with shape (nf, nall). + Lookup-safe atom types with shape (nf, nall). + real_atom + Physical-atom mask with shape (nf, nall). edge_index Edge source and destination indices with shape (2, E). edge_mask @@ -3266,7 +3273,12 @@ def _make_inter_potential_edge_mask( src = edge_index[0] dst = edge_index[1] atype_flat = atype.reshape(-1) - keep = edge_mask + real_atom_flat = real_atom.reshape(-1) + keep = ( + edge_mask + & real_atom_flat.index_select(0, src) + & real_atom_flat.index_select(0, dst) + ) descriptor = self.atomic_model.descriptor if descriptor.exclude_types: @@ -3274,9 +3286,7 @@ def _make_inter_potential_edge_mask( atom_excl = self.atomic_model.atom_excl if atom_excl is not None: - safe_atype, atom_is_present = self._sanitize_atom_types(atype) - atom_is_included = atom_is_present & atom_excl(safe_atype).to(torch.bool) - atom_is_included = atom_is_included.reshape(-1) + atom_is_included = atom_excl(atype).to(torch.bool).reshape(-1) keep = ( keep & atom_is_included.index_select(0, src) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 876a174e14..776bdece43 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -298,6 +298,11 @@ def get_dataloader_and_iter_lmdb( ) _block_targets = getattr(_data, "_block_targets", None) + _sampler_kwargs = { + "shuffle": True, + "seed": _training_params.get("seed"), + "block_targets": _block_targets, + } if self.world_size > 1: from deepmd.dpmodel.utils.lmdb_data import ( @@ -308,15 +313,12 @@ def get_dataloader_and_iter_lmdb( _data._reader, rank=self.rank, world_size=self.world_size, - shuffle=True, - seed=_training_params.get("seed"), - block_targets=_block_targets, + **_sampler_kwargs, ) else: _inner_sampler = LmdbBatchSampler( _data._reader, - shuffle=True, - block_targets=_block_targets, + **_sampler_kwargs, ) _dataloader = LmdbBatchDataLoader( From 38cc7b9a479ad4386a8369eadb6264d2bf7492cd Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 10 Aug 2026 22:16:03 +0800 Subject: [PATCH 4/9] fix(lmdb): preserve semantics across batch layouts --- deepmd/dpmodel/loss/ener.py | 88 ++++++++++++------- deepmd/dpmodel/loss/reduction.py | 21 +++-- deepmd/dpmodel/utils/lmdb_data.py | 19 ++-- deepmd/pt_expt/model/make_model.py | 41 ++++++--- deepmd/pt_expt/train/training.py | 13 +-- .../tests/common/dpmodel/test_loss_padding.py | 60 +++++++++++-- .../pt_expt/model/test_dpa2_graph_lower.py | 68 ++++++++++++-- source/tests/pt_expt/test_lmdb_training.py | 17 ++++ 8 files changed, 247 insertions(+), 80 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 1dd70361b7..e1679e3177 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -16,6 +16,12 @@ masked_pair_mean, per_frame_component_mean, ) +from deepmd.dpmodel.utils.neighbor_graph.graph import ( + frame_id_from_n_node, +) +from deepmd.dpmodel.utils.neighbor_graph.segment import ( + segment_sum, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -282,20 +288,34 @@ def call( # Two things about a batch decide how its terms reduce, and the node # axis states them differently. # - # ``inv``, the reciprocal real atom count of each frame, is what the - # extensive frame-level terms (energy, virial) divide by. ``maskf`` - # marks the padded rows the per-atom terms must skip. A rectangular - # batch carries both in its mask: summing it gives the counts, and its - # zeros are the padding. A ragged batch pads nothing, so it states the - # counts alone and its per-atom terms reduce over their whole axis. - maskf = None + # ``inv``, the reciprocal included-atom count of each frame, is what + # the extensive frame-level terms (energy, virial) divide by. + # ``maskf`` marks both padded rows and model-excluded atom types for + # the per-atom terms. A rectangular batch carries a two-dimensional + # mask whose row sums give the counts. A ragged batch carries the same + # information on its flat node axis, with ``n_node`` defining the + # frame segments. + maskf = ( + xp.astype(model_dict["mask"], energy.dtype) + if "mask" in model_dict + else None + ) + is_ragged = "n_node" in model_dict inv = None - if "n_node" in model_dict: - inv = 1.0 / xp.astype(model_dict["n_node"], energy.dtype) # [nf] - elif "mask" in model_dict: - maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] + if is_ragged: + n_node = model_dict["n_node"] + _nf = n_node.shape[0] + if maskf is None: + included_n_node = xp.astype(n_node, energy.dtype) + else: + frame_id = frame_id_from_n_node(n_node, n_total=maskf.shape[0]) + included_n_node = segment_sum(xp.reshape(maskf, (-1,)), frame_id, _nf) + inv = 1.0 / included_n_node + elif maskf is not None: inv = xp.reshape(1.0 / xp.sum(maskf, axis=-1), (-1,)) # [nf] _nloc = maskf.shape[1] + if maskf is not None: + _node_shape = maskf.shape if inv is not None: _nf = inv.shape[0] @@ -432,7 +452,7 @@ def call( l2_force_loss = xp.mean(xp.square(diff_f)) if maskf is not None: # Idiom 1 (per-atom masked mean, ncomp=3). - diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3] + diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3)) # Masked MSE computed for rmse_f display regardless of use_huber. l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3) if not self.use_huber: @@ -454,12 +474,12 @@ def call( ) # [nf, nloc, 3] huber_ncomp = 3 else: - diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3)) norm_2d = xp.reshape( xp.linalg.vector_norm( xp.reshape(diff_3, (-1, 3)), axis=1 ), - (_nf, _nloc), + _node_shape, ) abs_n = norm_2d quad_n = 0.5 * xp.square(norm_2d) @@ -470,7 +490,7 @@ def call( xp.where( abs_n <= self._huber_delta_force, quad_n, lin_n ), - (_nf, _nloc, 1), + (*_node_shape, 1), ) huber_ncomp = 1 l_huber_masked = masked_atom_mean( @@ -506,18 +526,18 @@ def call( ) elif self.loss_func == "mae": if maskf is not None: - diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3)) if not self.f_use_norm: l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: - diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3)) norm_2d = xp.reshape( xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), - (_nf, _nloc), + _node_shape, ) # One L2 norm per atom, hence one label per atom. l1_force_masked = masked_atom_mean( - xp.reshape(norm_2d, (_nf, _nloc, 1)), maskf, 1 + xp.reshape(norm_2d, (*_node_shape, 1)), maskf, 1 ) loss += pref_f * l1_force_masked more_loss["mae_f"] = self.display_if_exist( @@ -541,7 +561,7 @@ def call( ) if mae: if maskf is not None: - diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3)) mae_f = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: mae_f = xp.mean(xp.abs(diff_f)) @@ -627,10 +647,10 @@ def call( ) if maskf is not None: # Idiom 1 (per-atom masked mean, ncomp=1). - ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) - ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + ae_2d = xp.reshape(atom_ener, _node_shape) + ae_hat_2d = xp.reshape(atom_ener_hat, _node_shape) l2_ae_masked = masked_atom_mean( - xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + xp.square(ae_hat_2d - ae_2d)[..., None], maskf, 1 ) if not self.use_huber: loss += pref_ae * l2_ae_masked @@ -646,7 +666,7 @@ def call( abs_ae <= self._huber_delta_energy, quad_ae, lin_ae ) l_huber_ae_masked = masked_atom_mean( - huber_ae[:, :, None], maskf, 1 + huber_ae[..., None], maskf, 1 ) loss += pref_ae * l_huber_ae_masked more_loss["rmse_ae"] = self.display_if_exist( @@ -670,10 +690,10 @@ def call( xp.abs(atom_ener_hat_reshape - atom_ener_reshape) ) if maskf is not None: - ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) - ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + ae_2d = xp.reshape(atom_ener, _node_shape) + ae_hat_2d = xp.reshape(atom_ener_hat, _node_shape) l1_ae_masked = masked_atom_mean( - xp.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + xp.abs(ae_hat_2d - ae_2d)[..., None], maskf, 1 ) loss += pref_ae * l1_ae_masked more_loss["mae_ae"] = self.display_if_exist( @@ -697,8 +717,8 @@ def call( ) if maskf is not None: # Idiom 1 with pref weight (ncomp=3). - diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) - pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3)) + pf_3d = xp.reshape(atom_pref, (*_node_shape, 3)) l2_pf_masked = masked_atom_mean( xp.square(diff_f_3d) * pf_3d, maskf, 3 ) @@ -716,8 +736,8 @@ def call( xp.multiply(xp.abs(diff_f), atom_pref_reshape) ) if maskf is not None: - diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) - pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3)) + pf_3d = xp.reshape(atom_pref, (*_node_shape, 3)) l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3) loss += pref_pf * l1_pf_masked more_loss["mae_pf"] = self.display_if_exist( @@ -733,7 +753,7 @@ def call( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." ) if self.has_gf: - if maskf is None and inv is not None: + if is_ragged: # ``natoms`` below is one number for the whole batch, which a # padded batch can honour and a concatenated one cannot: its # frames differ in atom count, so ``drdq``, stored per frame @@ -787,6 +807,10 @@ def call( ) hessian = model_dict.get("hessian", model_dict.get("energy_derv_r_derv_r")) if self.has_h and hessian is not None and "hessian" in label_dict: + if is_ragged: + raise NotImplementedError( + "the hessian loss requires a rectangular atom axis" + ) find_hessian = label_dict.get("find_hessian", 0.0) if maskf is not None: hessian_shape = (_nf, _nloc * 3, _nloc * 3) diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py index 573b00684a..39980cc707 100644 --- a/deepmd/dpmodel/loss/reduction.py +++ b/deepmd/dpmodel/loss/reduction.py @@ -24,6 +24,10 @@ frame. Pooling and averaging over frames coincide there, so :func:`per_frame_component_mean` reduces per frame and leaves the frame axis to the caller, which applies the extensive ``1 / natoms`` weighting. +- **Pair terms** (Hessian) describe one response matrix per frame. Their + component count grows quadratically with atom count, so + :func:`masked_pair_mean` normalizes each matrix before averaging frames; + otherwise a large structure would dominate a batch quadratically. Pooling is what keeps a frame's weight independent of the company it keeps. The alternative -- averaging each frame's own per-label mean -- gives every @@ -80,11 +84,14 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: Parameters ---------- elem : Array - Non-negative per-element contribution of shape ``[nf, nloc, ncomp]`` - (already squared or abs'd, and pre-multiplied by any per-atom weight - such as ``atom_pref``). NOT yet multiplied by the mask. + Non-negative per-element contribution of shape + ``[*node_shape, ncomp]`` (already squared or abs'd, and pre-multiplied + by any per-atom weight such as ``atom_pref``). NOT yet multiplied by + the mask. ``node_shape`` may be rectangular ``(nf, nloc)`` or a flat + ragged node axis ``(N,)``. maskf : Array - Per-atom real/ghost mask of shape ``[nf, nloc]`` (1.0 real, 0.0 ghost). + Per-atom inclusion mask of shape ``node_shape`` (1.0 included, + 0.0 excluded). ncomp : int Number of components per atom (force: 3, atom energy: 1, dos: ``numb_dos``, tensor: ``tensor_size``). @@ -97,8 +104,10 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: ``0/0 = NaN``. """ xp = array_api_compat.array_namespace(elem, maskf) - total = xp.sum(elem * maskf[:, :, None]) - total_dof = xp.sum(maskf) * ncomp + node_elem = xp.reshape(elem, (-1, ncomp)) + node_mask = xp.reshape(maskf, (-1, 1)) + total = xp.sum(node_elem * node_mask) + total_dof = xp.sum(node_mask) * ncomp # A batch of nothing but padding has no label to average over, and the # ratio would be 0/0 = NaN -- poisoning the whole batch loss and, under # autograd, its gradient. The division still runs on a safe denominator so diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 888daf3166..dd3d79a276 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2327,12 +2327,15 @@ def ragged_batches(self) -> bool: return self._ragged_batches def use_ragged_batches(self, ragged: bool) -> None: - """Select the layout training batches are delivered in. - - The choice belongs to whichever model will consume them: one reading a - flat node axis takes the frames concatenated, one reading an - ``(nf, nloc, ...)`` axis needs them padded to a common width. Only the - trainer sees both the model and the data, so it makes the call, once, + """Select the layout mixed-nloc training batches are delivered in. + + The choice belongs to both the batching rule and the model. Only + ``mix:N`` may place different atom counts in one batch and therefore + needs a layout choice: a model reading a flat node axis takes those + frames concatenated, while one reading an ``(nf, nloc, ...)`` axis + needs them padded to a common width. Every other batching rule keeps + the established rectangular layout. Only the trainer sees both the + model and the data, so it requests the model-compatible layout once, before training starts. Consumers with a layout of their own -- statistics, validation -- name theirs at the point of use and are unaffected. @@ -2343,9 +2346,9 @@ def use_ragged_batches(self, ragged: bool) -> None: Parameters ---------- ragged : bool - Whether to concatenate frames instead of padding them. + Whether the consumer accepts concatenated mixed-nloc frames. """ - self._ragged_batches = ragged + self._ragged_batches = ragged and self.mixed_nloc def per_atom_strides(self) -> dict[str, int]: """Return the leading-axis entries per atom of each per-atom field. diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 00254e3efb..d51dd3648d 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -369,37 +369,52 @@ def _cal_hessian_ext_graph( """Graph twin of :func:`_cal_hessian_ext`. Computes the Hessian of the reduced output w.r.t. the LOCAL coordinates on - the carry-all graph route. Returns shape ``[nf, *vdef.shape, nloc*3, - nloc*3]`` -- the local-only counterpart of the dense extended Hessian, - already in the same final layout the dense route reaches after + the carry-all graph route. Each frame is differentiated only over its real + nodes; phantom rows are restored as zero rows and columns at the public + rectangular boundary. Returns shape ``[nf, *vdef.shape, nloc*3, nloc*3]`` + -- the local-only counterpart of the dense extended Hessian, already in + the same final layout the dense route reaches after ``communicate_extended_output`` folds ``nall -> nloc`` (the graph route reduces over owned nodes, so no fold is needed). Node axis is atom-major/xyz-minor, matching the dense final reshape. """ nf, nloc, _ = coord.shape vsize = math.prod(vdef.shape) - coord_flat = coord.reshape(nf, nloc * 3) hessians = [] for ii in range(nf): + node_index = torch.nonzero(atype[ii] >= 0, as_tuple=False).reshape(-1) + n_real = node_index.shape[0] + coord_flat = coord[ii, node_index].reshape(n_real * 3) + atype_frame = atype[ii : ii + 1, node_index] + aparam_frame = aparam[ii : ii + 1, node_index] if aparam is not None else None for ci in range(vsize): wrapper = _WrapperForwardEnergyGraph( model, kk, ci, - nloc, - atype[ii : ii + 1], + n_real, + atype_frame, box[ii : ii + 1] if box is not None else None, method, pair_excl, rcut, fparam[ii : ii + 1] if fparam is not None else None, - aparam[ii : ii + 1] if aparam is not None else None, + aparam_frame, ) hess = torch.autograd.functional.hessian( wrapper, - coord_flat[ii], + coord_flat, create_graph=create_graph, - ) # (nloc*3, nloc*3) + ) # (n_real*3, n_real*3) + if n_real != nloc: + component_index = ( + node_index[:, None] * 3 + + torch.arange(3, dtype=node_index.dtype, device=node_index.device) + ).reshape(-1) + hess = expand_node_values(hess, component_index, nloc * 3) + hess = expand_node_values( + hess.transpose(0, 1), component_index, nloc * 3 + ).transpose(0, 1) hessians.append(hess) return torch.stack(hessians).reshape(nf, *vdef.shape, nloc * 3, nloc * 3) @@ -947,9 +962,11 @@ def _call_common_graph( v, node_index, n_padded ).reshape(nf, nloc, *v.shape[1:]) # Graph-native Hessian (parallel to the dense ``forward_common_atomic`` - # loop): differentiate the reduced output w.r.t. the LOCAL coords by - # rebuilding the graph inside the wrapper. Added AFTER the unravel so - # its ``(nf, *def, nloc, 3, nloc, 3)`` shape is returned as-is. + # loop): differentiate the reduced output w.r.t. the compact LOCAL + # coordinates by rebuilding the graph inside the wrapper, then restore + # phantom rows and columns at the public rectangular boundary. Added + # AFTER the unravel so its ``(nf, *def, nloc, 3, nloc, 3)`` shape is + # returned as-is. # Eager-only, like the dense Hessian (autograd.functional.hessian # does not export/compile). aod = self.atomic_output_def() diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index e1769a117d..a954f5645c 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2199,12 +2199,13 @@ def _raise_if_full_validation_unsupported( def _configure_batch_layout(self, *data_maps: Any) -> None: """Ask each LMDB data system for the layout its own model can consume. - A model whose descriptor reads a flat node axis takes the frames of a - batch concatenated, which spares it the padding that frames of unequal - atom count would otherwise need. Every other model reads an - ``(nf, nloc, ...)`` axis and needs them padded to a common width. Only - the trainer sees both sides, and it settles the question here, once, - before any batch is drawn. + Under ``mix:N``, a model whose descriptor reads a flat node axis takes + the frames of a batch concatenated, which spares it the padding that + frames of unequal atom count would otherwise need. Every other model + reads an ``(nf, nloc, ...)`` axis and needs them padded to a common + width. Non-mixing batch rules retain their established rectangular + layout. Only the trainer sees both sides, and it settles the question + here, once, before any batch is drawn. A graph lower is necessary but not sufficient: the model must also expose an entry that takes the flat axis, which the composed models diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 1d51a6a3e1..98422193fa 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -568,35 +568,77 @@ def test_mask_and_n_node_agree(self, loss_func, intensive, use_huber): if use_huber and loss_func != "mse": pytest.skip("huber replaces the mse branch only") rng = np.random.default_rng(3) - nf, nloc = 3, 5 + nf, nloc = 2, 3 + n_node = np.array([2, 3], dtype=np.int64) + physical = np.arange(nloc)[None, :] < n_node[:, None] + mask = physical.astype(np.float64) + # The second physical node is excluded by the model, independently of + # the padding slot at the end of the first frame. + mask[0, 1] = 0.0 pred, label = _full_ener_dicts( nf, nloc, rng.normal(size=(nf, 1)), rng.normal(size=(nf, 1)), - mask=np.ones((nf, nloc), dtype=np.float64), + mask=mask, + force=rng.normal(size=(nf, nloc, 3)), virial=rng.normal(size=(nf, 9)), + atom_energy=rng.normal(size=(nf, nloc, 1)), + atom_ener=rng.normal(size=(nf, nloc, 1)), + atom_pref=rng.random(size=(nf, nloc * 3)), + find_force=1.0, find_virial=1.0, + find_atom_ener=1.0, + find_atom_pref=1.0, ) - loss_obj = self._loss( + label["force"] = rng.normal(size=(nf, nloc, 3)) + pref_pf = 0.0 if use_huber else 1.0 + loss_obj = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=1.0, + limit_pref_v=1.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + start_pref_pf=pref_pf, + limit_pref_pf=pref_pf, loss_func=loss_func, intensive_ener_virial=intensive, use_huber=use_huber, ) from_mask, more_mask = loss_obj.call(1.0, nloc, pred, label) - # The same batch, described by its counts instead of by a mask. - by_count = {k: v for k, v in pred.items() if k != "mask"} - by_count["n_node"] = np.full(nf, nloc, dtype=np.int64) - from_counts, more_counts = loss_obj.call(1.0, nloc, by_count, label) + # The same physical nodes on a flat ragged axis. ``n_node`` retains the + # frame boundaries while the flat mask retains the model exclusion. + by_count = { + **pred, + "force": pred["force"][physical], + "atom_energy": pred["atom_energy"][physical], + "mask": mask[physical], + "n_node": n_node, + } + ragged_label = { + **label, + "force": label["force"][physical], + "atom_ener": label["atom_ener"][physical], + "atom_pref": label["atom_pref"].reshape(nf, nloc, 3)[physical], + } + from_counts, more_counts = loss_obj.call( + 1.0, int(n_node.sum()), by_count, ragged_label + ) - np.testing.assert_allclose(float(from_counts), float(from_mask), rtol=0, atol=0) + np.testing.assert_allclose( + float(from_counts), float(from_mask), rtol=1e-14, atol=0 + ) assert sorted(more_counts) == sorted(more_mask) for key, value in more_mask.items(): np.testing.assert_allclose( np.asarray(more_counts[key], dtype=np.float64), np.asarray(value, dtype=np.float64), - rtol=0, + rtol=1e-14, atol=0, err_msg=key, ) diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py index 6e1a08481d..9fbdfa13ed 100644 --- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -200,23 +200,34 @@ def _make_model( # 1. routing/parity: proves the Task-7 generic plumbing works for DPA2. # ------------------------------------------------------------------ def test_force_virial_parity_vs_legacy(self) -> None: - """Default-flip ``forward_common`` (graph) matches - ``neighbor_graph_method="legacy"`` (dense) bit-tight at non-binding - sel with repformer attention off. + """The compact graph path matches dense and unpadded evaluation. + + A padded input exercises the graph path's compact-gather-scatter + round trip. At non-binding ``sel`` with repformer attention disabled, + it must match both the legacy dense route on the padded input and the + graph route on the equivalent input with the phantom suffix removed. """ model = self._make_model() # non-binding sel, attention off model.eval() box = self.cell.reshape(1, 9) + padded_atype = self.atype.clone() + padded_atype[:, -2:] = -1 + nreal = self.natoms - 2 graph = model.forward_common( - self.coord.clone().requires_grad_(True), self.atype, box + self.coord.clone().requires_grad_(True), padded_atype, box ) legacy = model.forward_common( self.coord.clone().requires_grad_(True), - self.atype, + padded_atype, box, neighbor_graph_method="legacy", ) + unpadded = model.forward_common( + self.coord[:, :nreal].clone().requires_grad_(True), + self.atype[:, :nreal], + box, + ) tol = {"rtol": 1e-10, "atol": 1e-10} torch.testing.assert_close(graph["energy_redu"], legacy["energy_redu"], **tol) torch.testing.assert_close( @@ -225,6 +236,21 @@ def test_force_virial_parity_vs_legacy(self) -> None: torch.testing.assert_close( graph["energy_derv_c_redu"], legacy["energy_derv_c_redu"], **tol ) + torch.testing.assert_close(graph["energy_redu"], unpadded["energy_redu"], **tol) + torch.testing.assert_close( + graph["energy_derv_r"][:, :nreal], + unpadded["energy_derv_r"], + **tol, + ) + torch.testing.assert_close( + graph["energy_derv_r"][:, nreal:], + torch.zeros_like(graph["energy_derv_r"][:, nreal:]), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph["energy_derv_c_redu"], unpadded["energy_derv_c_redu"], **tol + ) def test_graph_native_hessian_matches_dense(self) -> None: """The graph route computes the Hessian natively and it matches the @@ -255,15 +281,43 @@ def test_graph_native_hessian_matches_dense(self) -> None: assert _model_uses_graph_lower(model) is True box = self.cell.reshape(1, 9) + padded_atype = self.atype.clone() + padded_atype[:, -2:] = -1 + nreal = self.natoms - 2 + nreal_coord = nreal * 3 # ``EnergyModel.forward`` exposes the Hessian under the ``"hessian"`` # key (translated from ``energy_derv_r_derv_r``); it would KeyError if # the graph route failed to produce it. graph_out = model.forward( - self.coord.clone().requires_grad_(True), self.atype, box=box + self.coord.clone().requires_grad_(True), padded_atype, box=box ) assert "hessian" in graph_out assert graph_out["hessian"].shape == (1, self.natoms * 3, self.natoms * 3) + unpadded_out = model.forward( + self.coord[:, :nreal].clone().requires_grad_(True), + self.atype[:, :nreal], + box=box, + ) + torch.testing.assert_close( + graph_out["hessian"][:, :nreal_coord, :nreal_coord], + unpadded_out["hessian"], + rtol=1e-9, + atol=1e-9, + ) + torch.testing.assert_close( + graph_out["hessian"][:, nreal_coord:], + torch.zeros_like(graph_out["hessian"][:, nreal_coord:]), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_out["hessian"][:, :, nreal_coord:], + torch.zeros_like(graph_out["hessian"][:, :, nreal_coord:]), + rtol=0, + atol=0, + ) + # dense reference on the SAME weights via the escape hatch. ref_model = self._make_model() ref_model.eval() @@ -272,7 +326,7 @@ def test_graph_native_hessian_matches_dense(self) -> None: ref_model.enable_hessian() assert _model_uses_graph_lower(ref_model) is False dense_out = ref_model.forward( - self.coord.clone().requires_grad_(True), self.atype, box=box + self.coord.clone().requires_grad_(True), padded_atype, box=box ) torch.testing.assert_close( graph_out["hessian"], dense_out["hessian"], rtol=1e-9, atol=1e-9 diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index ab48c8b339..00213920fe 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -873,6 +873,23 @@ def test_graph_model_trains_on_a_flat_node_axis(self) -> None: # Frame-level fields keep their frame axis. self.assertEqual(batch["energy"].shape[0], batch["n_node"].shape[0]) + def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None: + """Model capability alone does not change established LMDB batches.""" + config = self._config(self._dpa1()) + config["training"]["training_data"]["batch_size"] = 2 + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + trainer = get_trainer(config) + batch = trainer.training_data.get_batch() + finally: + os.chdir(cwd) + + self.assertFalse(trainer.training_data._reader.ragged_batches) + self.assertEqual(batch["coord"].ndim, 3) + self.assertEqual(batch["atype"].ndim, 2) + self.assertNotIn("n_node", batch) + def test_compiled_graph_model_trains_on_a_flat_node_axis(self) -> None: """The compiled lower reads the flat axis too, so compiling changes nothing. From dc4c99a39cc406422682b095b2832fd09e0cdb71 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 11 Aug 2026 12:15:17 +0800 Subject: [PATCH 5/9] fix(lmdb): enforce mixed-batch model and loss contracts --- deepmd/dpmodel/loss/ener.py | 30 +++++- deepmd/dpmodel/loss/loss.py | 5 + deepmd/dpmodel/utils/lmdb_data.py | 22 +++-- deepmd/pt_expt/model/make_model.py | 91 +++++++++++++------ deepmd/pt_expt/train/training.py | 16 ++-- source/tests/common/dpmodel/test_lmdb_data.py | 14 +++ source/tests/common/dpmodel/test_loss_ener.py | 26 ++++++ .../tests/common/dpmodel/test_loss_padding.py | 27 ++++++ .../pt_expt/model/test_dpa2_graph_lower.py | 15 +++ .../pt_expt/model/test_dpa4_native_spin.py | 18 ++++ .../pt_expt/model/test_get_model_dpa4.py | 38 ++++++++ source/tests/pt_expt/test_lmdb_training.py | 25 +++++ 12 files changed, 278 insertions(+), 49 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index e1679e3177..8459d0847f 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -240,6 +240,11 @@ def __init__( if self.use_huber and self.has_h: raise RuntimeError("Huber loss is not implemented for hessian.") + @property + def supports_ragged_batches(self) -> bool: + """Whether the configured terms accept a flat per-node batch axis.""" + return not (self.has_gf or self.has_h) + def call( self, learning_rate: float, @@ -301,6 +306,8 @@ def call( else None ) is_ragged = "n_node" in model_dict + frame_id = None + included_n_node = None inv = None if is_ragged: n_node = model_dict["n_node"] @@ -310,10 +317,21 @@ def call( else: frame_id = frame_id_from_n_node(n_node, n_total=maskf.shape[0]) included_n_node = segment_sum(xp.reshape(maskf, (-1,)), frame_id, _nf) - inv = 1.0 / included_n_node elif maskf is not None: - inv = xp.reshape(1.0 / xp.sum(maskf, axis=-1), (-1,)) # [nf] + included_n_node = xp.reshape(xp.sum(maskf, axis=-1), (-1,)) _nloc = maskf.shape[1] + if included_n_node is not None: + has_included_node = included_n_node > 0 + safe_included_n_node = xp.where( + has_included_node, + included_n_node, + xp.ones_like(included_n_node), + ) + inv = xp.where( + has_included_node, + 1.0 / safe_included_n_node, + xp.zeros_like(included_n_node), + ) if maskf is not None: _node_shape = maskf.shape if inv is not None: @@ -331,7 +349,13 @@ def call( atom_ener = model_dict["atom_energy"] atom_ener_coeff = label_dict["atom_ener_coeff"] atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener.shape) - energy = xp.sum(atom_ener_coeff * atom_ener, axis=1) + weighted_atom_ener = atom_ener_coeff * atom_ener + if is_ragged: + if frame_id is None: + frame_id = frame_id_from_n_node(n_node, n_total=atom_ener.shape[0]) + energy = segment_sum(weighted_atom_ener, frame_id, _nf) + else: + energy = xp.sum(weighted_atom_ener, axis=1) if force_required: force_reshape = xp.reshape(force, (-1,)) force_hat_reshape = xp.reshape(force_hat, (-1,)) diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py index 3f45b12b8c..efee9c8c55 100644 --- a/deepmd/dpmodel/loss/loss.py +++ b/deepmd/dpmodel/loss/loss.py @@ -51,6 +51,11 @@ def call( def label_requirement(self) -> list[DataRequirementItem]: """Return data label requirements needed for this loss calculation.""" + @property + def supports_ragged_batches(self) -> bool: + """Whether this objective accepts a flat per-node batch axis.""" + return False + @staticmethod def display_if_exist(loss: Array, find_property: float) -> Array: """Display NaN if labeled property is not found. diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index dd3d79a276..0396640a4a 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -288,17 +288,19 @@ def _decode_frame( return result +def _canonical_lmdb_key(key: str) -> str: + """Return one LMDB field name in the in-memory DeePMD convention.""" + key = _KEY_REMAP.get(key, key) + if key.startswith("find_atomic_"): + return "find_atom_" + key.removeprefix("find_atomic_") + if key.startswith("atomic_"): + return "atom_" + key.removeprefix("atomic_") + return key + + def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]: """Remap LMDB key names to the canonical in-memory DeePMD convention.""" - out = {} - for k, v in frame.items(): - key = _KEY_REMAP.get(k, k) - if key.startswith("find_atomic_"): - key = "find_atom_" + key.removeprefix("find_atomic_") - elif key.startswith("atomic_"): - key = "atom_" + key.removeprefix("atomic_") - out[key] = v - return out + return {_canonical_lmdb_key(key): value for key, value in frame.items()} def _requirement_value(requirement: Any, key: str, default: Any) -> Any: @@ -381,7 +383,7 @@ def _raw_frame_availability( explicit_known = 0 explicit_true = 0 for key, value in frame.items(): - name = _KEY_REMAP.get(key, key) + name = _canonical_lmdb_key(key) if name.startswith("find_"): label = name.removeprefix("find_") bit = key_bits.get(label) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index d51dd3648d..f932b8847d 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -317,7 +317,9 @@ def __init__( pair_excl: Any, rcut: float, fparam: torch.Tensor | None, # (1, ndf) or None - aparam: torch.Tensor | None, # (1, nloc, nda) or None + aparam: torch.Tensor | None, # (nloc, nda) or None + spin: torch.Tensor | None, # (nloc, 3) or None + charge_spin: torch.Tensor | None, # (1, 2) or None ) -> None: self.model = model self.kk = kk @@ -330,6 +332,8 @@ def __init__( self.rcut = rcut self.fparam = fparam self.aparam = aparam + self.spin = spin + self.charge_spin = charge_spin def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: cc = coord_flat.reshape(1, self.nloc, 3) @@ -340,11 +344,9 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ng, self.atype.reshape(-1), fparam=self.fparam, - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) - # (N == nloc for the single-frame carry-all graph). - aparam=( - self.aparam.reshape(self.nloc, -1) if self.aparam is not None else None - ), + aparam=self.aparam, + spin=self.spin, + charge_spin=self.charge_spin, ) # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all # graph (all nodes owned); reduced output = sum over the node axis. @@ -361,6 +363,8 @@ def _cal_hessian_ext_graph( box: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + spin: torch.Tensor | None, + charge_spin: torch.Tensor | None, method: str, pair_excl: Any, rcut: float, @@ -380,26 +384,42 @@ def _cal_hessian_ext_graph( """ nf, nloc, _ = coord.shape vsize = math.prod(vdef.shape) + aparam_by_node = aparam.reshape(nf, nloc, -1) if aparam is not None else None + spin_by_node = spin.reshape(nf, nloc, 3) if spin is not None else None + charge_spin_by_frame = ( + charge_spin.reshape(1, -1) + if charge_spin is not None and charge_spin.ndim == 1 + else charge_spin + ) hessians = [] for ii in range(nf): node_index = torch.nonzero(atype[ii] >= 0, as_tuple=False).reshape(-1) n_real = node_index.shape[0] coord_flat = coord[ii, node_index].reshape(n_real * 3) atype_frame = atype[ii : ii + 1, node_index] - aparam_frame = aparam[ii : ii + 1, node_index] if aparam is not None else None + aparam_frame = ( + aparam_by_node[ii, node_index] if aparam_by_node is not None else None + ) + spin_frame = spin_by_node[ii, node_index] if spin_by_node is not None else None + charge_spin_frame = None + if charge_spin_by_frame is not None: + frame_index = 0 if charge_spin_by_frame.shape[0] == 1 else ii + charge_spin_frame = charge_spin_by_frame[frame_index : frame_index + 1] for ci in range(vsize): wrapper = _WrapperForwardEnergyGraph( - model, - kk, - ci, - n_real, - atype_frame, - box[ii : ii + 1] if box is not None else None, - method, - pair_excl, - rcut, - fparam[ii : ii + 1] if fparam is not None else None, - aparam_frame, + model=model, + kk=kk, + ci=ci, + nloc=n_real, + atype=atype_frame, + box=box[ii : ii + 1] if box is not None else None, + method=method, + pair_excl=pair_excl, + rcut=rcut, + fparam=fparam[ii : ii + 1] if fparam is not None else None, + aparam=aparam_frame, + spin=spin_frame, + charge_spin=charge_spin_frame, ) hess = torch.autograd.functional.hessian( wrapper, @@ -780,7 +800,8 @@ def call_common_ragged( Raises ------ NotImplementedError - If the model has no graph lower to read a flat node axis with. + If the model has no graph lower to read a flat node axis with, + or if its outputs require a rectangular pair axis. """ if not (self.mixed_types() and self.atomic_model.uses_graph_lower()): raise NotImplementedError( @@ -788,6 +809,14 @@ def call_common_ragged( "graph lower; this model reads a rectangular one, so its " "batches must be padded to a common atom count" ) + if any( + vdef.reducible and vdef.r_hessian + for vdef in self.atomic_output_def().get_data().values() + ): + raise NotImplementedError( + "Hessian outputs require a rectangular atom axis; a flat " + "node axis cannot represent their per-frame pair dimensions" + ) # The trainer resolves ``auto`` once and installs the concrete # builder on the model. A model reached outside it has none, and # resolving against its own device is what keeps that case from @@ -974,17 +1003,19 @@ def _call_common_graph( vdef = aod[kk] if vdef.reducible and vdef.r_hessian: model_predict[get_hessian_name(kk)] = _cal_hessian_ext_graph( - self, - kk, - vdef, - cc, - atype, - bb, - fp, - ap, - method, - pair_excl, - rcut, + model=self, + kk=kk, + vdef=vdef, + coord=cc, + atype=atype, + box=bb, + fparam=fp, + aparam=ap, + spin=spin, + charge_spin=charge_spin, + method=method, + pair_excl=pair_excl, + rcut=rcut, create_graph=self.training, ) return model_predict diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index a954f5645c..4460e585ce 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2207,12 +2207,14 @@ def _configure_batch_layout(self, *data_maps: Any) -> None: layout. Only the trainer sees both sides, and it settles the question here, once, before any batch is drawn. - A graph lower is necessary but not sufficient: the model must also - expose an entry that takes the flat axis, which the composed models - (linear, ZBL bridging) do not. Native-spin models also stay rectangular: - their public output translation needs the spin-specific force and mask, - while this generic ragged entry translates energy-model outputs only. - Requiring these capabilities keeps each model on the layout it can read. + A graph lower is necessary but not sufficient: the model must expose an + entry that takes the flat axis, and the configured loss must accept that + representation. Composed models (linear, ZBL bridging) have no such + model entry. Generalized-force and Hessian losses require rectangular + label axes. Native-spin models also stay rectangular because their + public output translation needs the spin-specific force and mask. + Requiring both capabilities keeps the complete objective on a layout it + can evaluate. Each task has its own data system and its own model, so the answer is each task's own: a multi-task run pairing a graph model with a dense @@ -2232,10 +2234,12 @@ def _configure_batch_layout(self, *data_maps: Any) -> None: """ for task_key in self.model_keys: model = self.models[task_key] + loss = self.losses[task_key] ragged = ( not model.has_spin() and model_uses_graph_lower(model) and hasattr(model, "forward_ragged") + and loss.supports_ragged_batches ) for data_map in data_maps: data = ( diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index ad8d5cf670..16dd1da1d1 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -722,6 +722,20 @@ def test_batch_demotes_mixed_label_availability(self): self.assertEqual(sorted(collated), sorted(batch)) np.testing.assert_allclose(collated["custom"], batch["custom"]) + def test_availability_scan_normalizes_atomic_label_prefix(self): + """Availability and frame decoding share the atomic-label aliases.""" + raw = msgpack.packb( + { + "atomic_dipole": np.ones(3).tolist(), + "find_atomic_dipole": 1.0, + }, + use_bin_type=True, + ) + + availability = lmdb_data_module._raw_frame_availability(raw, {"atom_dipole": 1}) + + self.assertEqual(availability, 1) + def test_availability_probe_is_deferred_until_requirements_are_registered(self): """Reader construction must not inspect labels before their use is known. diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py index b9dee7fe1e..d5c605fb0b 100644 --- a/source/tests/common/dpmodel/test_loss_ener.py +++ b/source/tests/common/dpmodel/test_loss_ener.py @@ -154,6 +154,32 @@ def test_forward(self) -> None: loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict) self.assertIsNotNone(loss) + def test_ragged_frames_are_reduced_independently(self) -> None: + """Atomic coefficients retain the frame boundaries of a flat node axis.""" + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + enable_atom_ener_coeff=True, + ) + n_node = np.array([2, 3], dtype=np.int64) + model_dict = { + "energy": np.array([[3.0], [12.0]]), + "atom_energy": np.arange(1.0, 6.0).reshape(-1, 1), + "n_node": n_node, + } + label_dict = { + "energy": np.array([[3.0], [12.0]]), + "find_energy": 1.0, + "atom_ener_coeff": np.ones((5, 1)), + } + + loss, _ = loss_fn.call(1.0, int(n_node.sum()), model_dict, label_dict) + + self.assertEqual(float(loss), 0.0) + class TestEnergyLossGeneralizedForce(TestEnergyLossBase): """Test energy loss with generalized force (numb_generalized_coord > 0). diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 98422193fa..92171b6f14 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -643,6 +643,33 @@ def test_mask_and_n_node_agree(self, loss_func, intensive, use_huber): err_msg=key, ) + @pytest.mark.parametrize("ragged", [False, True]) + def test_fully_excluded_frame_is_neutral(self, ragged): + """A frame without an included atom contributes no extensive label.""" + loss_obj = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + ) + mask_shape = (1,) if ragged else (1, 1) + model = { + "energy": np.ones((1, 1), dtype=np.float64), + "mask": np.zeros(mask_shape, dtype=np.float64), + } + if ragged: + model["n_node"] = np.ones(1, dtype=np.int64) + label = { + "energy": np.zeros((1, 1), dtype=np.float64), + "find_energy": 1.0, + } + + loss, more_loss = loss_obj.call(1.0, 1, model, label) + + assert float(loss) == 0.0 + assert float(more_loss["rmse_e"]) == 0.0 + def test_generalized_force_refuses_a_concatenated_batch(self): """``drdq`` is stored against a common atom axis, which is gone.""" rng = np.random.default_rng(5) diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py index 9fbdfa13ed..57e9df3514 100644 --- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -335,6 +335,21 @@ def test_graph_native_hessian_matches_dense(self) -> None: h = graph_out["hessian"][0] torch.testing.assert_close(h, h.transpose(-1, -2), rtol=1e-9, atol=1e-9) + def test_hessian_model_rejects_a_ragged_axis(self) -> None: + """A flat node axis cannot encode per-frame Hessian pair dimensions.""" + model = self._make_model().eval() + model.enable_hessian() + box = self.cell.reshape(1, 9) + n_node = torch.tensor([self.natoms], dtype=torch.int64, device=self.device) + + with pytest.raises(NotImplementedError, match="rectangular atom axis"): + model.forward_ragged( + self.coord.reshape(-1, 3).requires_grad_(True), + self.atype.reshape(-1), + n_node, + box=box, + ) + def test_disable_graph_lower_escape_hatch(self) -> None: """``descriptor.disable_graph_lower()`` is the documented legacy-dense escape hatch: it flips ``uses_graph_lower()`` to ``False`` so the diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index 865cf027f3..071dd40409 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -190,6 +190,24 @@ def _energy(sp: np.ndarray) -> float: atol=1e-6, ) + def test_hessian_uses_spin_conditioning(self) -> None: + """Coordinate Hessians retain the spin input of the energy surface.""" + self.model.enable_hessian() + hessian = self.model.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + self.box, + spin=self.spin, + )["energy_derv_r_derv_r"] + scaled_hessian = self.model.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + self.box, + spin=2.0 * self.spin, + )["energy_derv_r_derv_r"] + + assert torch.max(torch.abs(scaled_hessian - hessian)).item() > 1e-6 + def test_force_unchanged_by_spin_leaf_wiring(self) -> None: """``call_common`` WITHOUT ``spin`` has no ``energy_derv_r_mag`` key, and the spin-less forward is deterministic (the new ``spin is not diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py index aa76fd7ecc..a64d77e5eb 100644 --- a/source/tests/pt_expt/model/test_get_model_dpa4.py +++ b/source/tests/pt_expt/model/test_get_model_dpa4.py @@ -94,6 +94,44 @@ def test_get_model_normalized_config(self) -> None: self.assertEqual(ret["energy"].shape, (1, 1)) self.assertEqual(ret["force"].shape, (1, 5, 3)) + def test_graph_hessian_preserves_charge_spin(self) -> None: + """The Hessian is evaluated on the explicitly conditioned energy surface.""" + config = _make_raw_model_config() + config["descriptor"]["add_chg_spin_ebd"] = True + model = get_model(config).to(self.device).eval() + model.enable_hessian() + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], + dtype=torch.float64, + device=self.device, + requires_grad=True, + ) + atype = torch.tensor([[0, 1]], dtype=torch.int64, device=self.device) + box = (5.0 * torch.eye(3, dtype=torch.float64, device=self.device)).reshape( + 1, 9 + ) + charge_spin = torch.tensor( + [[0.0, 1.0]], dtype=torch.float64, device=self.device + ) + + graph = model.call_common(coord, atype, box, charge_spin=charge_spin) + dense = model.call_common( + coord, + atype, + box, + charge_spin=charge_spin, + neighbor_graph_method="legacy", + ) + + graph_hessian = graph["energy_derv_r_derv_r"] + self.assertEqual(graph_hessian.shape, (1, 1, 6, 6)) + torch.testing.assert_close( + graph_hessian, + dense["energy_derv_r_derv_r"], + rtol=1e-10, + atol=1e-10, + ) + def test_get_model_type_aliases(self) -> None: """All model-type aliases route to the SeZM path.""" for alias in ("dpa4", "DPA4", "sezm", "SeZM"): diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 00213920fe..65d0a50fce 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -873,6 +873,31 @@ def test_graph_model_trains_on_a_flat_node_axis(self) -> None: # Frame-level fields keep their frame axis. self.assertEqual(batch["energy"].shape[0], batch["n_node"].shape[0]) + def test_loss_contract_can_require_rectangular_batches(self) -> None: + """Generalized-force and Hessian labels keep their rectangular axes.""" + cases = { + "generalized force": { + "start_pref_gf": 1.0, + "limit_pref_gf": 1.0, + "numb_generalized_coord": 2, + }, + "Hessian": { + "start_pref_h": 1.0, + "limit_pref_h": 1.0, + }, + } + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + for name, loss_params in cases.items(): + with self.subTest(loss=name): + config = self._config(self._dpa1()) + config["loss"].update(loss_params) + trainer = get_trainer(config) + self.assertFalse(trainer.training_data._reader.ragged_batches) + finally: + os.chdir(cwd) + def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None: """Model capability alone does not change established LMDB batches.""" config = self._config(self._dpa1()) From c8c3a67a4df9f8c266d9ba32601d8da239f38869 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 12 Aug 2026 22:03:37 +0800 Subject: [PATCH 6/9] fix(loss): preserve boolean magnetic masks --- deepmd/dpmodel/loss/ener_spin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 042b4830b7..8015def9cd 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -328,7 +328,7 @@ def reshape_atomic(value: Array, ncomp: int) -> Array: pref_fm = pref_fm * find_force_mag force_mag_pred = reshape_atomic(model_dict["force_mag"], 3) force_mag_label = reshape_atomic(label_dict["force_mag"], 3) - mask_mag = reshape_atomic(model_dict["mask_mag"], 1) > 0 + mask_mag = reshape_atomic(model_dict["mask_mag"], 1) if maskf is not None: mask_mag = xp.logical_and( mask_mag, From d3b1b61940d8064cdd2eef94d71c1fcc894a2fad Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 13 Aug 2026 15:33:56 +0800 Subject: [PATCH 7/9] fix(pt): initialize legacy SeZM edge atom types --- deepmd/pt/model/model/sezm_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index de2fb58d58..5d9d182d79 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2849,7 +2849,7 @@ def build_edge_list_from_nlist( 1:1 with ``edge_index`` and ``edge_vec``. """ nloc = nlist.shape[1] - atype = torch.empty( + atype = torch.zeros( (nlist.shape[0], nloc), dtype=torch.long, device=extended_coord.device, From 0538fb88d6294ed003d15fecb51e09e1f9fcdbc5 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 13 Aug 2026 23:14:12 +0800 Subject: [PATCH 8/9] fix(pt_expt): handle virtual-spin statistics on CUDA --- deepmd/dpmodel/model/spin_model.py | 5 ++++- source/tests/pt_expt/test_lmdb_training.py | 5 +++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index afb805a89c..fc348c7730 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -20,6 +20,7 @@ ) from deepmd.dpmodel.common import ( NativeOP, + to_numpy_array, ) from deepmd.dpmodel.model.base_model import ( BaseModel, @@ -138,7 +139,9 @@ def __init__( self.spin_mask = self.spin.get_spin_mask() def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any: - """Convert a numpy array to the same namespace as ref_arr.""" + """Convert an array to the namespace and device of ``ref_arr``.""" + if array_api_compat.is_numpy_namespace(xp): + arr = to_numpy_array(arr) return xp.asarray(arr, device=array_api_compat.device(ref_arr)) def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array: diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 03eb5b2abd..84c053279a 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -54,6 +54,10 @@ DataRequirementItem, ) +from .compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + def _encode_array(arr: np.ndarray) -> dict: return { @@ -1016,6 +1020,7 @@ def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None: self.assertEqual(batch["atype"].ndim, 2) self.assertNotIn("n_node", batch) + @REQUIRES_SUPPORTED_COMPILE def test_compiled_graph_model_trains_on_a_flat_node_axis(self) -> None: """The compiled lower reads the flat axis too, so compiling changes nothing. From 8b41231d559441c407a39c700f34bf37a9f911cc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 14 Aug 2026 10:24:36 +0800 Subject: [PATCH 9/9] fix(dpmodel): handle cross-namespace device arrays --- deepmd/dpmodel/array_api.py | 40 ++++++++++++------- deepmd/dpmodel/descriptor/dpa4.py | 2 +- deepmd/dpmodel/model/spin_model.py | 10 +++-- deepmd/dpmodel/utils/exclude_mask.py | 19 +++++++-- source/tests/consistent/test_array_api.py | 35 ++++++++++++++++ source/tests/pt_expt/descriptor/test_dpa4.py | 24 +++++++++++ .../pt_expt/utils/test_exclusion_mask.py | 33 +++++++++++++++ 7 files changed, 140 insertions(+), 23 deletions(-) diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 0db1c51603..7f60aaa200 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -11,6 +11,10 @@ Version, ) +from deepmd.dpmodel.common import ( + to_numpy_array, +) + # Type alias for array_api compatible arrays Array = np.ndarray | Any # Any to support JAX, PyTorch, etc. arrays @@ -27,22 +31,28 @@ def xp_asarray_nodetach( ``torch.asarray`` detaches its input from the autograd graph, so calling ``xp.asarray`` on a weight attribute that is already a backend tensor (e.g. a ``torch.nn.Parameter`` registered by the pt_expt backend) - silently breaks gradient flow to that weight. This helper converts - genuine non-backend data (numpy arrays, python scalars/lists) via - ``xp.asarray``; backend tensors are returned as-is, with an optional - differentiable dtype cast via ``xp.astype``. - - The ``device`` argument only applies to the conversion path: backend - tensors are assumed to already live on the working device (they are - created together with the inputs). + silently breaks gradient flow to that weight. Backend tensors already in + ``xp`` are therefore returned as-is, with an optional differentiable dtype + cast via ``xp.astype``. + + An array from another namespace cannot retain its autograd graph. It is + converted through NumPy before entering ``xp``; this also performs the + required device-to-host copy when a CUDA-backed model constant is consumed + by a NumPy statistics path. + + The ``device`` argument only applies to the conversion path. Arrays already + in ``xp`` are assumed to live on the working device because model buffers + and inputs are moved together. """ - if isinstance(obj, np.ndarray) or not array_api_compat.is_array_api_obj(obj): - if dtype is None: - return xp.asarray(obj, device=device) - return xp.asarray(obj, dtype=dtype, device=device) - if dtype is not None and obj.dtype != dtype: - obj = xp.astype(obj, dtype) - return obj + if array_api_compat.is_array_api_obj(obj): + if array_api_compat.array_namespace(obj) is xp: + if dtype is not None and obj.dtype != dtype: + obj = xp.astype(obj, dtype) + return obj + obj = to_numpy_array(obj) + if dtype is None: + return xp.asarray(obj, device=device) + return xp.asarray(obj, dtype=dtype, device=device) # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816 diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index f1a08123c4..d96e23f84c 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2186,7 +2186,7 @@ def _canonicalize_charge_spin( raise ValueError("`charge_spin` is required for this SeZM descriptor.") charge_spin = xp.reshape( xp_asarray_nodetach( - xp, np.asarray(self.default_chg_spin), dtype=dtype, device=device + xp, self.default_chg_spin, dtype=dtype, device=device ), (1, 2), ) diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index fc348c7730..c7daebef71 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -14,13 +14,13 @@ from deepmd.dpmodel.array_api import ( Array, + xp_asarray_nodetach, ) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) from deepmd.dpmodel.common import ( NativeOP, - to_numpy_array, ) from deepmd.dpmodel.model.base_model import ( BaseModel, @@ -140,9 +140,11 @@ def __init__( def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any: """Convert an array to the namespace and device of ``ref_arr``.""" - if array_api_compat.is_numpy_namespace(xp): - arr = to_numpy_array(arr) - return xp.asarray(arr, device=array_api_compat.device(ref_arr)) + return xp_asarray_nodetach( + xp, + arr, + device=array_api_compat.device(ref_arr), + ) def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array: """Gather per-type values while mapping virtual atom types to zero. diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py index 1cc4ea7479..849bd2c2e9 100644 --- a/deepmd/dpmodel/utils/exclude_mask.py +++ b/deepmd/dpmodel/utils/exclude_mask.py @@ -5,6 +5,7 @@ from deepmd.dpmodel.array_api import ( Array, + xp_asarray_nodetach, xp_take_along_axis, xp_take_first_n, ) @@ -55,7 +56,11 @@ def build_type_exclude_mask( lead = atype.shape # (nf, natom) dense | (N,) graph return xp.reshape( xp.take( - xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)), + xp_asarray_nodetach( + xp, + self.type_mask[...], + device=array_api_compat.device(atype), + ), xp.reshape(atype, (-1,)), axis=0, ), @@ -151,7 +156,11 @@ def build_type_exclude_mask( type_ij_flat = xp.reshape(type_ij, (-1,)) mask = xp.reshape( xp.take( - xp.asarray(self.type_mask[...], device=array_api_compat.device(nlist)), + xp_asarray_nodetach( + xp, + self.type_mask[...], + device=array_api_compat.device(nlist), + ), type_ij_flat, ), (nf, nloc, nnei), @@ -185,7 +194,11 @@ def build_edge_exclude_mask(self, edge_index: Array, atype: Array) -> Array: dst_t = xp.take(atype, edge_index[1, :], axis=0) type_ij = dst_t * (self.ntypes + 1) + src_t return xp.take( - xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)), + xp_asarray_nodetach( + xp, + self.type_mask[...], + device=array_api_compat.device(atype), + ), type_ij, axis=0, ) diff --git a/source/tests/consistent/test_array_api.py b/source/tests/consistent/test_array_api.py index a9f2d8eaba..45a6cbe556 100644 --- a/source/tests/consistent/test_array_api.py +++ b/source/tests/consistent/test_array_api.py @@ -1,11 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import sys import unittest +from unittest.mock import ( + patch, +) +import array_api_compat import numpy as np from deepmd.dpmodel.array_api import ( xp_add_at, + xp_asarray_nodetach, xp_bincount, xp_maximum_at, xp_scatter_sum, @@ -56,6 +61,36 @@ def test_torch_parameter_requires_grad(self) -> None: self.assertTrue(param.requires_grad) self.assertEqual(param.device, DEVICE) + @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") + def test_foreign_tensor_is_converted_to_numpy_namespace(self) -> None: + tensor = torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE) + numpy_namespace = array_api_compat.array_namespace(np.empty(0)) + + # A CUDA tensor rejects the direct NumPy protocol. Simulate that + # boundary on CPU so the device-to-host fallback is exercised on every + # platform. + with patch.object( + torch.Tensor, + "numpy", + side_effect=TypeError("direct conversion is unavailable"), + ): + converted = xp_asarray_nodetach(numpy_namespace, tensor) + + self.assertIsInstance(converted, np.ndarray) + np.testing.assert_allclose(converted, np.array([1.0, 2.0])) + + @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") + def test_native_tensor_keeps_its_autograd_graph(self) -> None: + tensor = torch.nn.Parameter( + torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE) + ) + torch_namespace = array_api_compat.array_namespace(tensor) + + converted = xp_asarray_nodetach(torch_namespace, tensor) + + self.assertIs(converted, tensor) + self.assertTrue(converted.requires_grad) + class TestXpMaximumAtConsistent(unittest.TestCase): """Test maximum-at identities that differ between backend primitives.""" diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py index ed405be916..725fa642f4 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4.py +++ b/source/tests/pt_expt/descriptor/test_dpa4.py @@ -107,6 +107,30 @@ def test_consistency(self, use_env_seed, use_mapping) -> None: err_msg=err_msg, ) + def test_default_charge_spin_uses_model_namespace(self) -> None: + """A device buffer supplies the default without a NumPy round trip.""" + dtype = PRECISION_DICT["float64"] + descriptor = make_descriptor( + self.nt, + self.sel_mix, + self.rcut, + add_chg_spin_ebd=True, + default_chg_spin=[0.5, -0.5], + ).to(self.device) + coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) + atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) + nlist = torch.tensor(self.nlist, dtype=int, device=self.device) + + with mock.patch.object( + torch.Tensor, + "numpy", + side_effect=TypeError("direct conversion is unavailable"), + ): + output = descriptor(coord_ext, atype_ext, nlist)[0] + + assert output.device == self.device + assert torch.isfinite(output).all() + def test_train_and_eval_amp_switches_are_independent(self) -> None: """Training follows ``use_amp``, evaluation follows ``DP_AMP_INFER``. diff --git a/source/tests/pt_expt/utils/test_exclusion_mask.py b/source/tests/pt_expt/utils/test_exclusion_mask.py index cc0671c117..9ea89c12ac 100644 --- a/source/tests/pt_expt/utils/test_exclusion_mask.py +++ b/source/tests/pt_expt/utils/test_exclusion_mask.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from unittest.mock import ( + patch, +) import numpy as np import torch @@ -38,6 +41,13 @@ def test_build_type_exclude_mask(self) -> None: des = AtomExcludeMask(nt, exclude_types=exclude_types) mask = des.build_type_exclude_mask(torch.as_tensor(atype, device=env.DEVICE)) np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask) + with patch.object( + torch.Tensor, + "numpy", + side_effect=TypeError("direct conversion is unavailable"), + ): + numpy_mask = des.build_type_exclude_mask(atype) + np.testing.assert_equal(numpy_mask, expected_mask) def test_type_mask_is_buffer(self) -> None: des = AtomExcludeMask(3, exclude_types=[0]) @@ -66,6 +76,29 @@ def test_build_type_exclude_mask(self) -> None: torch.as_tensor(self.atype_ext, device=env.DEVICE), ) np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask) + with patch.object( + torch.Tensor, + "numpy", + side_effect=TypeError("direct conversion is unavailable"), + ): + numpy_mask = des.build_type_exclude_mask(self.nlist, self.atype_ext) + np.testing.assert_equal( + numpy_mask, + expected_mask, + ) + + def test_build_edge_exclude_mask_with_numpy_inputs(self) -> None: + des = PairExcludeMask(self.nt, exclude_types=[[0, 1]]) + edge_index = np.array([[0, 1, 2, 3], [1, 0, 3, 2]], dtype=np.int64) + atype = np.array([0, 1, 0, 0], dtype=np.int32) + + with patch.object( + torch.Tensor, + "numpy", + side_effect=TypeError("direct conversion is unavailable"), + ): + numpy_mask = des.build_edge_exclude_mask(edge_index, atype) + np.testing.assert_equal(numpy_mask, np.array([0, 0, 1, 1], dtype=np.int32)) def test_type_mask_is_buffer(self) -> None: des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])