Skip to content

Commit fb6ff93

Browse files
wanghan-iapcmHan Wang
andauthored
fix(cc): handle nloc==0 in DeepSpinPTExpt with phantom-atom padding (deepmodeling#5485)
## Problem Multi-rank spin MD can leave a rank with zero real local atoms (`nloc_real == 0`) when atoms migrate to other subdomains. The with-comm AOTI artifact hits an intermittent SIGFPE (integer divide by zero) at runtime in inductor-generated shape arithmetic that uses `nloc` as a divisor. Reproduced on master CI run [`26667802665`](https://github.com/deepmodeling/deepmd-kit/actions/runs/26667802665): ``` Caught signal 8 (Floating point exception: integer divide by zero) 4 forward_lower_with_comm/.../wrapper.so(AOTInductorModel::run_impl+0xf482) ``` Root cause: - The graph was traced with `nloc_min=1` (`serialization.py:362`) and inductor lowered an even stricter `nloc >= 2` runtime-check (visible in the generated `wrapper.cpp`'s `check_input_3`). - That runtime-check is gated by env var `AOTI_RUNTIME_CHECK_INPUTS` (default OFF), so with `nloc = 0` the check is silently bypassed and the compiled graph runs through its own divide-by-zero on shape arithmetic. - Whether the offending divide is actually emitted depends on inductor's code-gen choices, which vary across compiles — hence the intermittent nature. ## Fix Prepend two phantom atoms with empty neighbour lists when `nloc_real == 0` so the AOTI graph runs with `nloc == 2` and never reaches the integer-divide-by-zero path. Phantoms have no neighbours so they contribute zero atomic energy / force / virial, preserving the physically-correct "this rank has no real atoms" result. Key details (all in `source/api_cc/src/DeepSpinPTExpt.cc`): - `dcoord` / `datype` / `dspin` get two zero-valued rows prepended. - `firstneigh_tensor` gets two `-1` rows prepended (no neighbours). - `mapping_tensor` gets two identity entries prepended. - `comm_dict.nlocal` is set to `2` (not the LAMMPS-reported `0`) so `border_op` writes received ghost features past the phantom slots. - Output arrays (`dforce`, `dforce_mag`, `datom_energy`, `datom_virial`) get the phantom prefix stripped before being scattered back to LAMMPS via `select_map`. ## Why phantoms rather than `Dim(min=0)` re-export Bumping the trace constraint to `min=0` would require: 1. auditing every `nloc`-dependent divide in `deepmd/dpmodel/{descriptor,fitting,model}/` and protecting with `xp.maximum(nloc, 1)`; 2. `torch.export` re-emitting compatible guards (currently fails because spin-side shape relationships require `nloc >= 1` to be inferable); 3. inductor cooperating with the relaxed bound (it makes independent specialization choices downstream); 4. re-exporting every `.pt2` archive in `source/tests/infer/`. The phantom approach is a strict superset of correctness and self-contained in one C++ file. The two approaches aren't mutually exclusive — the `min=0` route can land as a follow-up once the dpmodel audit is done. ## Test plan - [x] Local CPU rebuild + `runUnitTests_cc --gtest_filter='*Spin*'`: **42 / 42 spin C++ regression tests pass** (12 TF-backend tests skipped, as expected in the PT-only venv). - [ ] CI: the multi-rank LAMMPS test `test_pair_deepmd_mpi_dpa3_spin_empty_subdomain` should now pass deterministically. Local Python LAMMPS-MPI verification is blocked by a pre-existing OpenMPI/MPICH ABI mismatch in my local venv (the plugin's `ompi_mpi_*` symbols can't resolve against MPICH's `libmpi.so.12`), so end-to-end verification falls to CI. ## Known limitations - The phantom path is structurally inert for `nloc_real > 0` (the `if (phantom_n > 0)` branch never fires), so the common path is unchanged. - If a future inductor version bumps the `nloc` lower-bound to >2, `phantom_n` will need to track that minimum. - This fix is in `DeepSpinPTExpt` only. The corresponding non-spin path in `DeepPotPTExpt` has the same code shape; non-spin DPA3 empty-subdomain currently passes in CI but could regress similarly with a future inductor change. Deferred to a follow-up if observed. - Supersedes deepmodeling#5478 (which proposed skipping the test); this PR fixes the underlying bug instead. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Prevented crashes and incorrect results when some processes have no local atoms during distributed runs; placeholder (phantom) atoms are ignored in neighbor, force, energy, and per-atom outputs so reported values match real atoms. * **Tests** * Updated tests to pass explicit per-atom parameters and exercise empty-subdomain and multi-rank behaviors. * **Documentation** * Clarified test docstrings and model config comments about per-atom parameter handling. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 65c0420 commit fb6ff93

5 files changed

Lines changed: 143 additions & 13 deletions

File tree

source/api_cc/src/DeepSpinPTExpt.cc

Lines changed: 106 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -371,12 +371,51 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
371371
int nloc = nall_real - nghost_real;
372372
int nframes = 1;
373373

374-
// Build spin tensor for real atoms using bkw_map
375-
std::vector<VALUETYPE> dspin(static_cast<size_t>(nall_real) * 3);
376-
for (int ii = 0; ii < nall_real; ++ii) {
374+
// Phantom-atom padding for the empty-subdomain corner case
375+
// (``nloc_real == 0``). Multi-rank spin MD can land a rank with zero
376+
// real local atoms when atoms migrate to other subdomains. The
377+
// with-comm AOTI artifact, traced with ``nloc_min=1`` and lowered by
378+
// inductor with an even stricter ``nloc >= 2`` runtime-check
379+
// (silently bypassed because ``AOTI_RUNTIME_CHECK_INPUTS`` is unset by
380+
// default), then SIGFPEs at runtime with an "integer divide by zero"
381+
// inside inductor-generated shape arithmetic that uses ``nloc`` as a
382+
// divisor. The failure is intermittent because inductor re-codegens
383+
// across runs and only some compiles emit the offending divide.
384+
//
385+
// Fix: prepend two phantom atoms with no neighbours so the AOTI graph
386+
// runs with ``nloc == 2``. The phantoms have an empty nlist row and
387+
// therefore contribute zero atomic energy / force / virial, preserving
388+
// the physically-correct "this rank has no real atoms" semantics.
389+
// ``nlocal`` in the comm tensors is set to ``2`` so border_op writes
390+
// received ghost features past the phantom slots; outputs are stripped
391+
// of the phantom prefix before being scattered back to LAMMPS atoms
392+
// via ``select_map``.
393+
const int phantom_n = (nloc_real == 0 && nall_real > 0) ? 2 : 0;
394+
if (phantom_n > 0) {
395+
dcoord.insert(dcoord.begin(), static_cast<size_t>(phantom_n) * 3,
396+
static_cast<VALUETYPE>(0));
397+
datype.insert(datype.begin(), static_cast<size_t>(phantom_n), 0);
398+
// Keep aparam_ aligned with the padded local atoms: the phantom atoms
399+
// get zero-valued atomic-parameter rows so the aparam tensor built below
400+
// (shape {1, nloc, daparam}) stays consistent with the padded ``nloc``.
401+
// (aparam_nall is false here, so aparam_ is a per-local-atom buffer.)
402+
if (daparam > 0) {
403+
aparam_.insert(aparam_.begin(), static_cast<size_t>(phantom_n) * daparam,
404+
static_cast<VALUETYPE>(0));
405+
}
406+
nall_real += phantom_n;
407+
nloc_real = phantom_n;
408+
nloc = nall_real - nghost_real;
409+
}
410+
411+
// Build spin tensor for real atoms using bkw_map (skip phantom prefix
412+
// which keeps zero spin).
413+
std::vector<VALUETYPE> dspin(static_cast<size_t>(nall_real) * 3,
414+
static_cast<VALUETYPE>(0));
415+
for (int ii = phantom_n; ii < nall_real; ++ii) {
377416
for (int dd = 0; dd < 3; ++dd) {
378417
dspin[static_cast<size_t>(ii) * 3 + dd] =
379-
spin[static_cast<size_t>(bkw_map[ii]) * 3 + dd];
418+
spin[static_cast<size_t>(bkw_map[ii - phantom_n]) * 3 + dd];
380419
}
381420
}
382421

@@ -445,11 +484,25 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
445484
nlist_data.shuffle_exclude_empty(fwd_map);
446485
nlist_data.padding();
447486

448-
// Rebuild mapping tensor
487+
// Rebuild mapping tensor. Phantom slots (when phantom_n > 0) get
488+
// identity entries — they index into their own row and never appear
489+
// in any other atom's nlist (their nlist rows are all -1 below).
449490
if (lmp_list.mapping) {
450491
std::vector<std::int64_t> mapping(nall_real);
451-
for (int ii = 0; ii < nall_real; ii++) {
452-
mapping[ii] = fwd_map[lmp_list.mapping[bkw_map[ii]]];
492+
for (int ii = 0; ii < phantom_n; ii++) {
493+
mapping[ii] = ii;
494+
}
495+
for (int ii = phantom_n; ii < nall_real; ii++) {
496+
// Defensive: this branch (lmp_list.mapping != nullptr) is single-rank
497+
// only (set_mapping is gated on comm->nprocs==1 in pair_deepspin /
498+
// pair_deepmd), while phantom_n>0 only occurs on a multi-rank empty
499+
// subdomain, so the two cannot currently co-occur and the +phantom_n
500+
// term is a no-op (phantom_n==0) on every reachable path. It is kept
501+
// so the mapping stays correct -- resolving fwd_map's pre-padding local
502+
// index into the post-padding local index space -- if that invariant
503+
// ever changes.
504+
mapping[ii] =
505+
fwd_map[lmp_list.mapping[bkw_map[ii - phantom_n]]] + phantom_n;
453506
}
454507
mapping_tensor =
455508
torch::from_blob(mapping.data(), {1, nall_real}, int_option)
@@ -472,8 +525,16 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
472525
}
473526

474527
// Flatten raw nlist — the .pt2 model sorts by distance on-device.
528+
// Phantom rows (all -1) are prepended below so the AOTI graph sees
529+
// nloc == phantom_n + nloc_real_orig instead of 0.
475530
firstneigh_tensor =
476531
createNlistTensor(nlist_data.jlist, nnei).to(torch::kInt64).to(device);
532+
if (phantom_n > 0) {
533+
auto phantom_rows = torch::full(
534+
{1, phantom_n, nnei}, static_cast<std::int64_t>(-1),
535+
torch::TensorOptions().dtype(torch::kInt64).device(device));
536+
firstneigh_tensor = torch::cat({phantom_rows, firstneigh_tensor}, 1);
537+
}
477538
}
478539

479540
// Build fparam/aparam tensors
@@ -566,6 +627,23 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
566627
ener.assign(flat_energy_.data_ptr<ENERGYTYPE>(),
567628
flat_energy_.data_ptr<ENERGYTYPE>() + flat_energy_.numel());
568629

630+
// Zero the reduced energy on an empty rank. Phantoms have constant
631+
// atomic outputs (per-type bias + zero-neighbour MLP) that flow into
632+
// ``energy_redu`` -- and on the spin path the SpinModel doubles atoms
633+
// so the bias contribution appears for both real and spin phantom
634+
// halves; subtracting only the real-half exposed by
635+
// ``output_map["energy"]`` after the ``[:, :nloc]`` slice leaves the
636+
// spin-half leaking into the MPI-reduced LAMMPS total. The physical
637+
// contribution of a rank with no real local atoms is zero by
638+
// definition, so just clear ``ener`` directly.
639+
//
640+
// Forces, force_mag, and virial are unaffected because phantom atomic
641+
// outputs are coord-independent (no neighbours) so their derivatives
642+
// are zero -- no analogous correction is needed.
643+
if (phantom_n > 0) {
644+
std::fill(ener.begin(), ener.end(), static_cast<ENERGYTYPE>(0));
645+
}
646+
569647
// Extract force: energy_derv_r (nf, nall, 1, 3) -> (nf, nall, 3)
570648
torch::Tensor force_tensor =
571649
output_map["energy_derv_r"].squeeze(-2).view({-1}).to(floatType);
@@ -588,6 +666,17 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
588666
virial.assign(cpu_virial_.data_ptr<VALUETYPE>(),
589667
cpu_virial_.data_ptr<VALUETYPE>() + cpu_virial_.numel());
590668

669+
// Strip the phantom prefix (see phantom-atom padding comment near
670+
// ``select_real_atoms_coord``) so the ``bkw_map`` lookup below sees
671+
// only the real / ghost atoms it was built for. The phantom slots
672+
// carry zero forces because their nlist rows were all -1 — they
673+
// produce no neighbour contributions, so dropping them is exact.
674+
if (phantom_n > 0) {
675+
dforce.erase(dforce.begin(), dforce.begin() + phantom_n * 3);
676+
dforce_mag.erase(dforce_mag.begin(), dforce_mag.begin() + phantom_n * 3);
677+
nall_real -= phantom_n;
678+
}
679+
591680
// bkw map: map force from real atoms back to full atom list
592681
force.resize(static_cast<size_t>(nframes) * fwd_map.size() * 3);
593682
force_mag.resize(static_cast<size_t>(nframes) * fwd_map.size() * 3);
@@ -612,6 +701,16 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener,
612701
cpu_atom_virial_.data_ptr<VALUETYPE>(),
613702
cpu_atom_virial_.data_ptr<VALUETYPE>() + cpu_atom_virial_.numel());
614703

704+
// Strip the phantom prefix from atomic outputs as well (see force
705+
// block above). Phantom slots carry zero atomic energy / virial
706+
// because their nlist rows were all -1.
707+
if (phantom_n > 0) {
708+
datom_energy.erase(datom_energy.begin(),
709+
datom_energy.begin() + phantom_n);
710+
datom_virial.erase(datom_virial.begin(),
711+
datom_virial.begin() + phantom_n * 9);
712+
}
713+
615714
atom_energy.resize(static_cast<size_t>(nframes) * fwd_map.size());
616715
atom_virial.resize(static_cast<size_t>(nframes) * fwd_map.size() * 9);
617716
select_map<VALUETYPE>(atom_energy, datom_energy, bkw_map, 1, nframes,

source/api_cc/tests/test_with_comm_load_failure_ptexpt.cc

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,8 +177,10 @@ TEST_F(TestDeepSpinPTExptWithCommLoadFailure, single_rank_compute_succeeds) {
177177

178178
double ener;
179179
std::vector<double> force_, force_mag, virial;
180+
// The fixture is built with numb_aparam=1; supply a uniform per-atom aparam.
181+
std::vector<double> fparam, aparam(natoms, 1.0);
180182
EXPECT_NO_THROW(dp.compute(ener, force_, force_mag, virial, coord, spin,
181-
atype, empty_box, 0, inlist, 0));
183+
atype, empty_box, 0, inlist, 0, fparam, aparam));
182184
}
183185

184186
TEST_F(TestDeepSpinPTExptWithCommLoadFailure, multi_rank_compute_throws) {
@@ -192,11 +194,17 @@ TEST_F(TestDeepSpinPTExptWithCommLoadFailure, multi_rank_compute_throws) {
192194
deepmd::InputNlist inlist(natoms, ilist.data(), numneigh.data(),
193195
firstneigh.data());
194196
convert_nlist(inlist, nlist_data);
195-
inlist.nswap = 1; // simulate multi-rank without populating send/recv
197+
// Multi-rank is keyed on nprocs (DeepSpinPTExpt.cc), not nswap; with
198+
// has_comm_artifact_ true but the with-comm loader failed to load, the
199+
// dispatch must throw.
200+
inlist.nprocs = 2;
196201

197202
double ener;
198203
std::vector<double> force_, force_mag, virial;
204+
// The fixture is built with numb_aparam=1; supply a uniform per-atom aparam
205+
// so the throw comes from the multi-rank dispatch, not a missing aparam.
206+
std::vector<double> fparam, aparam(natoms, 1.0);
199207
EXPECT_THROW(dp.compute(ener, force_, force_mag, virial, coord, spin, atype,
200-
empty_box, 0, inlist, 0),
208+
empty_box, 0, inlist, 0, fparam, aparam),
201209
deepmd::deepmd_exception);
202210
}

source/lmp/tests/run_mpi_pair_deepmd_spin_dpa3_pt2.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,12 @@
106106
lammps.timestep(0.0005)
107107
lammps.fix("1 all nve")
108108

109-
lammps.pair_style(f"deepspin {args.PB_FILE}")
109+
# The DPA3 spin fixture is built with numb_aparam=1, so supply a uniform
110+
# atom parameter. This exercises the aparam path in DeepSpinPTExpt, including
111+
# the empty-subdomain phantom-atom aparam padding; a uniform value keeps the
112+
# per-rank results self-consistent (real atoms get the same aparam regardless
113+
# of the processor grid).
114+
lammps.pair_style(f"deepspin {args.PB_FILE} aparam 1.0")
110115
lammps.pair_coeff(args.pair_coeff)
111116
lammps.compute("virial all centroid/stress/atom NULL pair")
112117
# Per-atom magnetic force components. LAMMPS does not expose ``fm``

source/lmp/tests/test_lammps_spin_dpa3_pt2.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,10 @@ def test_pair_deepmd_mpi_dpa3_spin_empty_subdomain() -> None:
263263
empty-rank guard for the spin path (the with-comm artifact still
264264
runs on rank 1 with nloc_real=0). Compares against same-archive
265265
mpi-1 reference.
266+
267+
The DPA3 spin fixture has ``numb_aparam=1`` and the runner supplies a
268+
uniform aparam, so the empty rank also exercises the phantom-atom aparam
269+
padding in ``DeepSpinPTExpt`` (PR #5485 review).
266270
"""
267271
out_mpi = _run_mpi_subprocess(nprocs=2, data_path=data_file_empty_subdomain)
268272
out_ref = _run_mpi_subprocess(nprocs=1, data_path=data_file_empty_subdomain)

source/tests/infer/gen_spin.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,14 @@ def _build_dpa3_mpi_yaml(yaml_path: str) -> None:
125125
"precision": "float64",
126126
"seed": 1,
127127
},
128-
"fitting_net": {"neuron": [5, 5, 5], "resnet_dt": True, "seed": 1},
128+
# numb_aparam=1 exercises the aparam path of DeepSpinPTExpt, including
129+
# the empty-subdomain phantom-atom aparam padding (PR #5485 review).
130+
"fitting_net": {
131+
"neuron": [5, 5, 5],
132+
"resnet_dt": True,
133+
"numb_aparam": 1,
134+
"seed": 1,
135+
},
129136
"spin": {"use_spin": [True, False], "virtual_scale": [0.3140, 0.0]},
130137
}
131138

@@ -185,7 +192,14 @@ def _build_dpa3_single_yaml(yaml_path: str) -> None:
185192
"precision": "float64",
186193
"seed": 1,
187194
},
188-
"fitting_net": {"neuron": [5, 5, 5], "resnet_dt": True, "seed": 1},
195+
# numb_aparam=1 exercises the aparam path of DeepSpinPTExpt, including
196+
# the empty-subdomain phantom-atom aparam padding (PR #5485 review).
197+
"fitting_net": {
198+
"neuron": [5, 5, 5],
199+
"resnet_dt": True,
200+
"numb_aparam": 1,
201+
"seed": 1,
202+
},
189203
"spin": {"use_spin": [True, False], "virtual_scale": [0.3140, 0.0]},
190204
}
191205

0 commit comments

Comments
 (0)