diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index 446b2386ec..f0a101d8b1 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -276,6 +276,59 @@ Run under Kokkos with one GPU: lmp -k on g 1 -sf kk -in in.lammps ``` +### Model-deviation sampling with Kokkos + +`deepmd/kk` supports model-deviation sampling with one MPI rank. The first +model supplies the energy, force, and virial that drive dynamics on every +step. Each model is loaded once. Additional models are evaluated only at the +existing `out_freq` interval, including step zero; `out_freq 0` disables +sampling. The first model's already computed outputs are reused, so a +two-model sample takes two model evaluations in total. + +```lammps +atom_modify map yes +pair_style deepmd/kk m0.pt2 m1.pt2 out_freq 1000 out_file model_devi.out +pair_coeff * * H C N O P Cl +``` + +Run with `-k on g 1 -sf kk -pk kokkos neigh full newton off`. Multi-model +sampling requires float32 compact canonical `.pt2` artifacts with identical +cutoffs and complete type maps. Models must have zero frame and atomic +parameter dimensions and no spin types. A compressed archive's fixed default +charge/spin state is supported: its API may report a nonzero charge/spin +width because an optional fold can reset that state, even though the compiled +canonical forward has no runtime charge/spin input. Reference models must +have matching charge/spin widths and identical frozen default states, including +older compressed archives without a charge-state fold. Every model must export +atomic virials, even when the `atomic` output option is disabled, because the +device path uses them to compute the total virial. Missing or incompatible +default states and missing atomic-virial outputs are rejected during model +initialization. +Explicit `charge_spin` settings and runtime parameters from `compute`, `fix`, +or `ttm` are unsupported in this mode. Incompatible models are rejected; +there is no fallback to host sampling. The single-rank periodic +box restriction still applies: its thickness must exceed twice the cutoff in +every periodic direction. Multi-model runs with more than one MPI rank are +rejected; existing single-model edge/graph and multi-rank execution is +unchanged. + +All models consume the same canonical graph on the GPU. Kokkos accumulates +force and virial deviations on the GPU with population statistics; no +coordinates, neighbor lists, raw forces, or raw virials are copied to the +host for sampling. Without `atomic`, only six final statistics and a status +flag are copied back, and the file contains seven columns including the +timestep. `atomic` optionally adds the final per-atom force deviations. +`relative`, `relative_v`, units, and restart output retain ordinary `deepmd` +semantics. The mean force deviation is the arithmetic mean over atoms; the +reported average virial deviation is the RMS over nine tensor components. + +Sampling runs after the first model's driving outputs have been consumed; +reference outputs never replace its dynamics. The existing graph builder's +node-map refresh and LAMMPS control/thermodynamic output remain unchanged. +The device-resident sampling claim concerns the additional model evaluations +and statistics, not every operation performed by LAMMPS. Include the extra +reference inference and device statistics when comparing performance. + ### CPU hosts DPA4C has a second set of hand-written operators for the CPU, so a compressed diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index d98d744387..9c26df7366 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -13,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 30 +#define DP_C_API_VERSION 31 /** * @brief Neighbor list. @@ -2614,6 +2614,34 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp); */ int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp); +/** + * @brief Copy the current default charge/spin state, including frozen states. + * Supported by the PyTorch Exportable backend. Zero values identify an + * unconditioned model; an unknown state is an error. + * Clears any previous error before querying; DP_DeepPotCheckOK reports only + * the result of this query. + * @param[in] dp An initialized DP. + * @param[out] values Caller-owned buffer, or NULL with capacity zero to query + * the required length. No values are written on error. + * @param[in] capacity Buffer capacity in doubles; must fit the entire state. + * @return Required length, or -1 on error. Inspect DP_DeepPotCheckOK for + * errors. The length is independent of DP_DeepPotGetDimChgSpin. + * @since API version 31 + */ +int DP_DeepPotGetDefaultChgSpin(DP_DeepPot* dp, double* values, int capacity); + +/** + * @brief Whether the model provides atomic virials. + * Supported by the PyTorch Exportable backend. + * Clears any previous error before querying; DP_DeepPotCheckOK reports only + * the result of this query, including successful queries returning false. + * @param[in] dp An initialized DP. + * @return True if available, false if unavailable or on error. Inspect + * DP_DeepPotCheckOK to distinguish an unsupported query from unavailability. + * @since API version 31 + */ +bool DP_DeepPotHasAtomicVirial(DP_DeepPot* dp); + /** * @brief Fix the charge/spin condition served for the rest of the run. * It becomes the condition of every later evaluation that is not given one diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 681b197b0e..c7bc90ab05 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1287,6 +1287,33 @@ class DeepPot : public DeepBaseModel { return dchgspin; } + /** + * @brief Get the current default charge/spin state, including frozen states. + * An empty vector identifies an unconditioned model. Throws for an unknown + * state or unsupported backend. The length need not equal dim_chg_spin(). + */ + std::vector get_default_chg_spin() const { + assert(dp); + const int size = DP_DeepPotGetDefaultChgSpin(dp, nullptr, 0); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + std::vector state(size); + if (size > 0) { + DP_DeepPotGetDefaultChgSpin(dp, state.data(), size); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + return state; + } + + /** + * @brief Whether the model provides atomic virials; throws if unsupported. + */ + bool has_atomic_virial() const { + assert(dp); + const bool result = DP_DeepPotHasAtomicVirial(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + /** * @brief Fix the charge/spin condition served for the rest of the run. * It becomes the condition of every later evaluation that is not given one diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index c35a4a1763..4e23aa4042 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -2790,6 +2790,44 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp) { int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); } +int DP_DeepPotGetDefaultChgSpin(DP_DeepPot* dp, double* values, int capacity) { + try { + dp->exception.clear(); + if (capacity < 0 || (values == nullptr && capacity != 0)) { + throw deepmd::deepmd_exception( + "default charge/spin output requires a nonnegative capacity and " + "a buffer, or NULL with capacity zero for a length query"); + } + const auto state = dp->dp.get_default_chg_spin(); + if (state.size() > + static_cast(std::numeric_limits::max())) { + throw deepmd::deepmd_exception("default charge/spin state is too large"); + } + const int size = static_cast(state.size()); + if (values != nullptr) { + if (capacity < size) { + throw deepmd::deepmd_exception( + "default charge/spin output buffer is too small"); + } + std::copy(state.begin(), state.end(), values); + } + return size; + } catch (const deepmd::deepmd_exception& ex) { + dp->exception = ex.what(); + return -1; + } +} + +bool DP_DeepPotHasAtomicVirial(DP_DeepPot* dp) { + try { + dp->exception.clear(); + return dp->dp.has_atomic_virial(); + } catch (const deepmd::deepmd_exception& ex) { + dp->exception = ex.what(); + return false; + } +} + void DP_DeepPotSetChargeSpin(DP_DeepPot* dp, const double* charge_spin, const int numb_chg_spin) { diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 575c7924c9..e867e0319f 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -208,6 +208,19 @@ class DeepPotBackend : public DeepBaseModelBackend { **/ virtual int dim_chg_spin() const { return 0; } + /** + * @brief Get the current default charge/spin state, including frozen states. + * An empty vector identifies an unconditioned model. Throws if the backend + * cannot establish the state. Its length need not equal dim_chg_spin(). + */ + virtual std::vector get_default_chg_spin() const; + + /** + * @brief Whether the loaded model provides atomic virials. + * Throws if the backend cannot query this capability. + */ + virtual bool has_atomic_virial() const; + /** * @brief Reject a charge/spin condition no embedding table can address. * @@ -909,6 +922,12 @@ class DeepPot : public DeepBaseModel { int dim_chg_spin() const; + /** @copydoc DeepPotBackend::get_default_chg_spin */ + std::vector get_default_chg_spin() const; + + /** @copydoc DeepPotBackend::has_atomic_virial */ + bool has_atomic_virial() const; + /** * @brief Fix the charge/spin condition served for the rest of the run. * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 4707ee9ce6..e7252a79f0 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -413,6 +413,8 @@ class DeepPotPTExpt : public DeepPotBackend { */ bool uses_fp32_edge_vectors() const override; bool uses_canonical_graph_inference() const override; + std::vector get_default_chg_spin() const override; + bool has_atomic_virial() const override; private: template @@ -460,6 +462,8 @@ class DeepPotPTExpt : public DeepPotBackend { bool has_default_fparam_; std::vector default_fparam_; std::vector default_chg_spin_; + // Defer unavailable metadata to the query, preserving single-model loading. + std::string default_chg_spin_query_error_; /** Half-open row range of each charge-state value, from the archive. */ std::vector > chg_spin_table_ranges_; double rcut; diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 3373d6fe71..c625743918 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -692,6 +692,16 @@ bool DeepPotBackend::supports_device_edge_inference() const { return false; } bool DeepPotBackend::uses_canonical_graph_inference() const { return false; } +std::vector DeepPotBackend::get_default_chg_spin() const { + throw deepmd::deepmd_exception( + "this backend cannot query the model's default charge/spin state"); +} + +bool DeepPotBackend::has_atomic_virial() const { + throw deepmd::deepmd_exception( + "this backend cannot query atomic-virial availability"); +} + void DeepPot::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -778,6 +788,12 @@ bool DeepPot::uses_canonical_graph_inference() const { int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } +std::vector DeepPot::get_default_chg_spin() const { + return dp->get_default_chg_spin(); +} + +bool DeepPot::has_atomic_virial() const { return dp->has_atomic_virial(); } + void DeepPot::set_charge_spin(const std::vector& charge_spin) { dp->set_charge_spin(charge_spin); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 005a18e722..b5d8ec819b 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -407,6 +407,15 @@ void DeepPotPTExpt::init(const std::string& model, check_charge_spin_domain(default_chg_spin_, chg_spin_table_ranges_); } + // A legacy compressed archive can carry a frozen default without shipping + // a fold. Retain that state without changing either conditioning width. + try { + default_chg_spin_ = + deepmd::ptexpt::read_default_chg_spin_for_query(metadata, dchgspin); + } catch (const deepmd::deepmd_exception& ex) { + default_chg_spin_query_error_ = ex.what(); + } + int num_intra_nthreads, num_inter_nthreads; get_env_nthreads(num_intra_nthreads, num_inter_nthreads); if (num_inter_nthreads) { @@ -451,6 +460,7 @@ void DeepPotPTExpt::set_charge_spin(const std::vector& charge_spin) { *loader); } default_chg_spin_.swap(next_charge_spin); + default_chg_spin_query_error_.clear(); } std::vector DeepPotPTExpt::run_model( @@ -2549,6 +2559,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a compact canonical artifact."); } + if (!has_atomic_virial()) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu requires atomic virials, but the model " + "was exported without them or does not declare atom_virial output."); + } if (!gpu_enabled) { throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); @@ -2700,4 +2715,21 @@ bool DeepPotPTExpt::uses_canonical_graph_inference() const { return lower_input_is_canonical_; } +std::vector DeepPotPTExpt::get_default_chg_spin() const { + assert(inited); + if (!default_chg_spin_query_error_.empty()) { + throw deepmd::deepmd_exception(default_chg_spin_query_error_); + } + return default_chg_spin_; +} + +bool DeepPotPTExpt::has_atomic_virial() const { + assert(inited); + // Canonical graphs expose public output names; other lower schemas use + // the internal energy derivative names consumed by compute(). + const char* key = lower_input_is_canonical_ ? "atom_virial" : "energy_derv_c"; + return do_atomic_virial && std::find(output_keys.begin(), output_keys.end(), + key) != output_keys.end(); +} + #endif diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index 8b8696f55f..6924fb03b4 100644 --- a/source/api_cc/src/commonPTExpt.h +++ b/source/api_cc/src/commonPTExpt.h @@ -381,6 +381,83 @@ inline void check_charge_spin_domain( } } +/** + * @brief Read the served charge/spin state independently of the input width. + * Compressed archives can freeze a state without a runtime input or a fold. + * An empty result requires explicit metadata identifying no conditioning. + */ +inline std::vector read_default_chg_spin_for_query( + const JsonValue& metadata, const int dim_chg_spin) { + for (const auto* key : {"has_chg_spin_ebd", "has_default_chg_spin"}) { + if (metadata.has(key) && metadata[key].type != JsonValue::Bool) { + throw deepmd::deepmd_exception(std::string(key) + " must be boolean."); + } + } + if (metadata.has("dim_chg_spin") && + (metadata["dim_chg_spin"].type != JsonValue::Number || + metadata["dim_chg_spin"].as_double() != dim_chg_spin || + dim_chg_spin < 0)) { + throw deepmd::deepmd_exception( + "dim_chg_spin must be a nonnegative integer."); + } + std::vector state; + if (metadata.has("default_chg_spin") && + metadata["default_chg_spin"].type != JsonValue::Null) { + const auto& value = metadata["default_chg_spin"]; + if (value.type != JsonValue::Array) { + throw deepmd::deepmd_exception("default_chg_spin must be an array."); + } + for (const auto& component : value.as_array()) { + if (component.type != JsonValue::Number || + !std::isfinite(component.as_double()) || + component.as_double() != std::floor(component.as_double())) { + throw deepmd::deepmd_exception( + "default_chg_spin must contain finite integer values."); + } + state.push_back(component.as_double()); + } + } + if (metadata.has("has_default_chg_spin") && + metadata["has_default_chg_spin"].as_bool() != !state.empty()) { + throw deepmd::deepmd_exception( + "has_default_chg_spin disagrees with default_chg_spin."); + } + const bool unconditioned = metadata.has("has_chg_spin_ebd") && + !metadata["has_chg_spin_ebd"].as_bool(); + if (unconditioned && (dim_chg_spin != 0 || !state.empty())) { + throw deepmd::deepmd_exception( + "has_chg_spin_ebd=false disagrees with the charge/spin state."); + } + if (state.empty()) { + if (unconditioned) { + return state; + } + throw deepmd::deepmd_exception( + "the model's default charge/spin state cannot be established from " + "metadata; regenerate the model with an explicit default state"); + } + if (dim_chg_spin > 0 && + state.size() != static_cast(dim_chg_spin)) { + throw deepmd::deepmd_exception( + "default_chg_spin length does not match dim_chg_spin."); + } + if (metadata.has("chg_spin_table_ranges")) { + for (const auto& bounds : metadata["chg_spin_table_ranges"].as_array()) { + for (const auto& limit : bounds.as_array()) { + if (limit.type != JsonValue::Number || + !std::isfinite(limit.as_double())) { + throw deepmd::deepmd_exception( + "chg_spin_table_ranges must contain finite numeric bounds."); + } + } + } + } + check_charge_spin_domain( + state, + read_chg_spin_table_ranges(metadata, static_cast(state.size()))); + return state; +} + /** * @brief Validate a charge/spin condition supplied with an inference call. * diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index a8e559d0ee..e184c83d06 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -7,6 +7,7 @@ #include #include #include +#include #include #include "DeepPot.h" @@ -125,6 +126,65 @@ TEST(TestPtExptMetadata, default_chg_spin_length_mismatch_throws) { deepmd::deepmd_exception); } +TEST(TestPtExptMetadata, query_retains_frozen_default_without_a_fold) { + const auto metadata = deepmd::ptexpt::parse_json( + R"({"dim_chg_spin": 0, "has_chg_spin_ebd": true, + "default_chg_spin": [0, 1]})"); + // The runtime-input reader must still report no input. The query instead + // reports the state frozen into this legacy artifact's constants. + EXPECT_TRUE(deepmd::ptexpt::read_default_chg_spin(metadata, 0).empty()); + EXPECT_EQ(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 0), + (std::vector{0.0, 1.0})); +} + +TEST(TestPtExptMetadata, query_identifies_unconditioned_models) { + const auto metadata = deepmd::ptexpt::parse_json( + R"({"dim_chg_spin": 0, "has_chg_spin_ebd": false, + "has_default_chg_spin": false, "default_chg_spin": null})"); + EXPECT_TRUE( + deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 0).empty()); +} + +TEST(TestPtExptMetadata, query_rejects_unknown_or_malformed_states) { + for (const auto* json : { + R"({})", + R"({"has_chg_spin_ebd": true})", + R"({"has_chg_spin_ebd": true, "default_chg_spin": []})", + R"({"has_chg_spin_ebd": false, "default_chg_spin": [0, 1]})", + R"({"has_default_chg_spin": true, "default_chg_spin": null})", + R"({"has_default_chg_spin": false, "default_chg_spin": [0, 1]})", + R"({"has_chg_spin_ebd": "false"})", + R"({"default_chg_spin": "0, 1"})", + R"({"default_chg_spin": [0, "1"]})", + R"({"default_chg_spin": [0, 1.5]})", + R"({"dim_chg_spin": "0", "default_chg_spin": [0, 1]})", + R"({"default_chg_spin": [0, 1], + "chg_spin_table_ranges": [["0", 100], [0, 100]]})", + }) { + SCOPED_TRACE(json); + const auto metadata = deepmd::ptexpt::parse_json(json); + EXPECT_THROW(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 0), + deepmd::deepmd_exception); + } +} + +TEST(TestPtExptMetadata, query_validates_input_width_and_state_domain) { + auto metadata = deepmd::ptexpt::parse_json( + R"({"default_chg_spin": [0, 1], + "chg_spin_table_ranges": [[-100, 100], [0, 100]]})"); + EXPECT_EQ(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 2), + (std::vector{0.0, 1.0})); + EXPECT_THROW(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 1), + deepmd::deepmd_exception); + metadata.obj_val["default_chg_spin"].arr_val[0].num_val = 100.0; + EXPECT_THROW(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 0), + deepmd::deepmd_exception); + metadata.obj_val["default_chg_spin"].arr_val[0].num_val = + std::numeric_limits::infinity(); + EXPECT_THROW(deepmd::ptexpt::read_default_chg_spin_for_query(metadata, 0), + deepmd::deepmd_exception); +} + // A compressed archive reports width zero at load and names its width only // once the charge-state fold is read, so the ranges are read twice. The first // read must yield nothing rather than judge the count it cannot yet know. @@ -177,6 +237,8 @@ TYPED_TEST(TestInferDeepPotAPtExpt, cpu_build_nlist) { deepmd::DeepPot& dp = this->dp; EXPECT_FALSE(dp.supports_device_edge_inference()); EXPECT_FALSE(dp.uses_canonical_graph_inference()); + EXPECT_TRUE(dp.get_default_chg_spin().empty()); + EXPECT_TRUE(dp.has_atomic_virial()); double ener; std::vector force, virial; dp.compute(ener, force, virial, coord, atype, box); @@ -694,6 +756,7 @@ TYPED_TEST(TestInferDeepPotAPtExpt, cpu_atomic_throws_when_disabled) { deepmd::DeepPot dp_no_av; ASSERT_NO_THROW( dp_no_av.init("../../tests/infer/deeppot_sea_no_atomic_virial.pt2")); + EXPECT_FALSE(dp_no_av.has_atomic_virial()); std::vector& coord = this->coord; std::vector& atype = this->atype; diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index a03482a90a..cd78549c4f 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -478,8 +478,7 @@ void PairDeepMD::compute(int eflag, int vflag) { scale[1][1] * dvatom[9 * ii + 5] * ener_unit_cvt_factor; // zy } } - if (out_freq > 0 && update->ntimestep % out_freq == 0) { - int rank = comm->me; + if (model_deviation_step()) { // std force if (newton_pair) { #if LAMMPS_VERSION_NUMBER >= 20220324 @@ -488,102 +487,7 @@ void PairDeepMD::compute(int eflag, int vflag) { comm->reverse_comm_pair(this); #endif } - vector std_f; - vector tmp_avg_f; - deep_pot_model_devi.compute_avg(tmp_avg_f, all_force); - deep_pot_model_devi.compute_std_f(std_f, tmp_avg_f, all_force); - if (out_rel == 1) { - deep_pot_model_devi.compute_relative_std_f(std_f, tmp_avg_f, eps); - } - double min = numeric_limits::max(), max = 0, avg = 0; - ana_st(max, min, avg, std_f, nlocal); - double all_f_min = 0, all_f_max = 0, all_f_avg = 0; - MPI_Reduce(&min, &all_f_min, 1, MPI_DOUBLE, MPI_MIN, 0, world); - MPI_Reduce(&max, &all_f_max, 1, MPI_DOUBLE, MPI_MAX, 0, world); - MPI_Reduce(&avg, &all_f_avg, 1, MPI_DOUBLE, MPI_SUM, 0, world); - all_f_avg /= double(atom->natoms); - // std v - std::vector send_v(9 * numb_models); - std::vector recv_v(9 * numb_models); - for (int kk = 0; kk < numb_models; ++kk) { - for (int ii = 0; ii < 9; ++ii) { - send_v[kk * 9 + ii] = all_virial[kk][ii] / double(atom->natoms); - } - } - MPI_Reduce(&send_v[0], &recv_v[0], 9 * numb_models, MPI_DOUBLE, MPI_SUM, - 0, world); - std::vector> all_virial_1(numb_models); - std::vector avg_virial, std_virial; - for (int kk = 0; kk < numb_models; ++kk) { - all_virial_1[kk].resize(9); - for (int ii = 0; ii < 9; ++ii) { - all_virial_1[kk][ii] = recv_v[kk * 9 + ii]; - } - } - double all_v_min = numeric_limits::max(), all_v_max = 0, - all_v_avg = 0; - if (rank == 0) { - deep_pot_model_devi.compute_avg(avg_virial, all_virial_1); - deep_pot_model_devi.compute_std(std_virial, avg_virial, all_virial_1, - 1); - if (out_rel_v == 1) { - deep_pot_model_devi.compute_relative_std(std_virial, avg_virial, - eps_v, 1); - } - for (int ii = 0; ii < 9; ++ii) { - if (std_virial[ii] > all_v_max) { - all_v_max = std_virial[ii]; - } - if (std_virial[ii] < all_v_min) { - all_v_min = std_virial[ii]; - } - all_v_avg += std_virial[ii] * std_virial[ii]; - } - all_v_avg = sqrt(all_v_avg / 9); - } - if (rank == 0) { - all_v_max *= ener_unit_cvt_factor; - all_v_min *= ener_unit_cvt_factor; - all_v_avg *= ener_unit_cvt_factor; - all_f_max *= force_unit_cvt_factor; - all_f_min *= force_unit_cvt_factor; - all_f_avg *= force_unit_cvt_factor; - fp << setw(12) << update->ntimestep << " " << setw(18) << all_v_max - << " " << setw(18) << all_v_min << " " << setw(18) << all_v_avg - << " " << setw(18) << all_f_max << " " << setw(18) << all_f_min - << " " << setw(18) << all_f_avg; - } - if (out_each == 1) { - vector std_f_all(atom->natoms); - // Gather std_f and tags - tagint* tag = atom->tag; - int nprocs = comm->nprocs; - ensure_model_deviation_buffers(); - for (int ii = 0; ii < nlocal; ii++) { - tagsend[ii] = tag[ii]; - stdfsend[ii] = std_f[ii]; - } - MPI_Gather(&nlocal, 1, MPI_INT, counts, 1, MPI_INT, 0, world); - displacements[0] = 0; - for (int ii = 0; ii < nprocs - 1; ii++) { - displacements[ii + 1] = displacements[ii] + counts[ii]; - } - MPI_Gatherv(tagsend, nlocal, MPI_LMP_TAGINT, tagrecv, counts, - displacements, MPI_LMP_TAGINT, 0, world); - MPI_Gatherv(stdfsend, nlocal, MPI_DOUBLE, stdfrecv, counts, - displacements, MPI_DOUBLE, 0, world); - if (rank == 0) { - for (int dd = 0; dd < atom->natoms; ++dd) { - std_f_all[tagrecv[dd] - 1] = stdfrecv[dd] * force_unit_cvt_factor; - } - for (int dd = 0; dd < atom->natoms; ++dd) { - fp << " " << setw(18) << std_f_all[dd]; - } - } - } - if (rank == 0) { - fp << endl; - } + write_model_deviation(all_virial); } } else { error->all(FLERR, "unknown computational branch"); @@ -621,6 +525,124 @@ void PairDeepMD::compute(int eflag, int vflag) { } } +bool PairDeepMD::model_deviation_step() const { + return numb_models > 1 && out_freq > 0 && update->ntimestep % out_freq == 0; +} + +void PairDeepMD::write_model_deviation( + const std::vector>& all_virial) { + int nlocal = atom->nlocal; + const int rank = comm->me; + vector std_f; + vector tmp_avg_f; + deep_pot_model_devi.compute_avg(tmp_avg_f, all_force); + deep_pot_model_devi.compute_std_f(std_f, tmp_avg_f, all_force); + if (out_rel == 1) { + deep_pot_model_devi.compute_relative_std_f(std_f, tmp_avg_f, eps); + } + double min = numeric_limits::max(), max = 0, avg = 0; + ana_st(max, min, avg, std_f, nlocal); + double all_f_min = 0, all_f_max = 0, all_f_avg = 0; + MPI_Reduce(&min, &all_f_min, 1, MPI_DOUBLE, MPI_MIN, 0, world); + MPI_Reduce(&max, &all_f_max, 1, MPI_DOUBLE, MPI_MAX, 0, world); + MPI_Reduce(&avg, &all_f_avg, 1, MPI_DOUBLE, MPI_SUM, 0, world); + all_f_avg /= double(atom->natoms); + // std v + std::vector send_v(9 * numb_models); + std::vector recv_v(9 * numb_models); + for (int kk = 0; kk < numb_models; ++kk) { + for (int ii = 0; ii < 9; ++ii) { + send_v[kk * 9 + ii] = all_virial[kk][ii] / double(atom->natoms); + } + } + MPI_Reduce(&send_v[0], &recv_v[0], 9 * numb_models, MPI_DOUBLE, MPI_SUM, 0, + world); + std::vector> all_virial_1(numb_models); + std::vector avg_virial, std_virial; + for (int kk = 0; kk < numb_models; ++kk) { + all_virial_1[kk].resize(9); + for (int ii = 0; ii < 9; ++ii) { + all_virial_1[kk][ii] = recv_v[kk * 9 + ii]; + } + } + double all_v_min = numeric_limits::max(), all_v_max = 0, + all_v_avg = 0; + if (rank == 0) { + deep_pot_model_devi.compute_avg(avg_virial, all_virial_1); + deep_pot_model_devi.compute_std(std_virial, avg_virial, all_virial_1, 1); + if (out_rel_v == 1) { + deep_pot_model_devi.compute_relative_std(std_virial, avg_virial, eps_v, + 1); + } + for (int ii = 0; ii < 9; ++ii) { + if (std_virial[ii] > all_v_max) { + all_v_max = std_virial[ii]; + } + if (std_virial[ii] < all_v_min) { + all_v_min = std_virial[ii]; + } + all_v_avg += std_virial[ii] * std_virial[ii]; + } + all_v_avg = sqrt(all_v_avg / 9); + } + write_model_deviation_output( + {all_v_max, all_v_min, all_v_avg, all_f_max, all_f_min, all_f_avg}, + std_f); +} + +void PairDeepMD::write_model_deviation_output( + const std::array& deviation, const std::vector& std_f) { + const int rank = comm->me; + int nlocal = atom->nlocal; + double all_v_max = deviation[0], all_v_min = deviation[1], + all_v_avg = deviation[2]; + double all_f_max = deviation[3], all_f_min = deviation[4], + all_f_avg = deviation[5]; + if (rank == 0) { + all_v_max *= ener_unit_cvt_factor; + all_v_min *= ener_unit_cvt_factor; + all_v_avg *= ener_unit_cvt_factor; + all_f_max *= force_unit_cvt_factor; + all_f_min *= force_unit_cvt_factor; + all_f_avg *= force_unit_cvt_factor; + fp << setw(12) << update->ntimestep << " " << setw(18) << all_v_max << " " + << setw(18) << all_v_min << " " << setw(18) << all_v_avg << " " + << setw(18) << all_f_max << " " << setw(18) << all_f_min << " " + << setw(18) << all_f_avg; + } + if (out_each == 1) { + vector std_f_all(atom->natoms); + // Gather std_f and tags + tagint* tag = atom->tag; + int nprocs = comm->nprocs; + ensure_model_deviation_buffers(); + for (int ii = 0; ii < nlocal; ii++) { + tagsend[ii] = tag[ii]; + stdfsend[ii] = std_f[ii]; + } + MPI_Gather(&nlocal, 1, MPI_INT, counts, 1, MPI_INT, 0, world); + displacements[0] = 0; + for (int ii = 0; ii < nprocs - 1; ii++) { + displacements[ii + 1] = displacements[ii] + counts[ii]; + } + MPI_Gatherv(tagsend, nlocal, MPI_LMP_TAGINT, tagrecv, counts, displacements, + MPI_LMP_TAGINT, 0, world); + MPI_Gatherv(stdfsend, nlocal, MPI_DOUBLE, stdfrecv, counts, displacements, + MPI_DOUBLE, 0, world); + if (rank == 0) { + for (int dd = 0; dd < atom->natoms; ++dd) { + std_f_all[tagrecv[dd] - 1] = stdfrecv[dd] * force_unit_cvt_factor; + } + for (int dd = 0; dd < atom->natoms; ++dd) { + fp << " " << setw(18) << std_f_all[dd]; + } + } + } + if (rank == 0) { + fp << endl; + } +} + static bool is_key(const string& input) { vector keys; keys.push_back("out_freq"); @@ -646,6 +668,15 @@ static bool is_key(const string& input) { return false; } +bool PairDeepMD::initialize_models(const std::vector& models) { + deep_pot.init(models[0], get_node_rank(), get_file_content(models[0])); + if (models.size() > 1) { + deep_pot_model_devi.init(models, get_node_rank(), get_file_content(models)); + return true; + } + return false; +} + void PairDeepMD::settings(int narg, char** arg) { if (narg <= 0) { error->all(FLERR, "Illegal pair_style command"); @@ -663,39 +694,23 @@ void PairDeepMD::settings(int narg, char** arg) { models.push_back(arg[ii]); } numb_models = models.size(); - if (numb_models == 1) { - try { - deep_pot.init(arg[0], get_node_rank(), get_file_content(arg[0])); - } catch (deepmd_compat::deepmd_exception& e) { - error->one(FLERR, e.what()); - } - cutoff = deep_pot.cutoff() * dist_unit_cvt_factor; - numb_types = deep_pot.numb_types(); - numb_types_spin = deep_pot.numb_types_spin(); - dim_fparam = deep_pot.dim_fparam(); - dim_aparam = deep_pot.dim_aparam(); - dim_chg_spin = deep_pot.dim_chg_spin(); - } else { - try { - deep_pot.init(arg[0], get_node_rank(), get_file_content(arg[0])); - deep_pot_model_devi.init(models, get_node_rank(), - get_file_content(models)); - } catch (deepmd_compat::deepmd_exception& e) { - error->one(FLERR, e.what()); - } - cutoff = deep_pot_model_devi.cutoff() * dist_unit_cvt_factor; - numb_types = deep_pot_model_devi.numb_types(); - numb_types_spin = deep_pot_model_devi.numb_types_spin(); - dim_fparam = deep_pot_model_devi.dim_fparam(); - dim_aparam = deep_pot_model_devi.dim_aparam(); - dim_chg_spin = deep_pot_model_devi.dim_chg_spin(); - assert(cutoff == deep_pot.cutoff() * dist_unit_cvt_factor); - assert(numb_types == deep_pot.numb_types()); - assert(numb_types_spin == deep_pot.numb_types_spin()); - assert(dim_fparam == deep_pot.dim_fparam()); - assert(dim_aparam == deep_pot.dim_aparam()); - assert(dim_chg_spin == deep_pot.dim_chg_spin()); + if (numb_models == 0) { + error->all(FLERR, "No model specified for pair style deepmd."); + } + bool ensemble_initialized = false; + try { + ensemble_initialized = initialize_models(models); + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); } + // Ensemble metadata have always been taken from model 0. Device styles can + // therefore use the same parser without constructing a host ensemble. + cutoff = deep_pot.cutoff() * dist_unit_cvt_factor; + numb_types = deep_pot.numb_types(); + numb_types_spin = deep_pot.numb_types_spin(); + dim_fparam = deep_pot.dim_fparam(); + dim_aparam = deep_pot.dim_aparam(); + dim_chg_spin = deep_pot.dim_chg_spin(); out_freq = 100; out_file = "model_devi.out"; @@ -881,7 +896,7 @@ void PairDeepMD::settings(int narg, char** arg) { if (!charge_spin.empty()) { try { deep_pot.set_charge_spin(charge_spin); - if (numb_models > 1) { + if (ensemble_initialized) { deep_pot_model_devi.set_charge_spin(charge_spin); } } catch (deepmd_compat::deepmd_exception& e) { diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index c7c40d9b18..5aecbce05f 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -28,6 +28,7 @@ namespace deepmd_compat = deepmd; namespace deepmd_compat = deepmd::hpp; #endif +#include #include #include #include @@ -54,10 +55,21 @@ class PairDeepMD : public PairDeepBaseModel { protected: deepmd_compat::DeepPot deep_pot; deepmd_compat::DeepPotModelDevi deep_pot_model_devi; + /** Load models and return whether the legacy ensemble accepts settings. */ + virtual bool initialize_models(const std::vector& models); // Assemble the send/recv swap metadata (a comm-only neighbor list; its // geometry fields are unused) for the device-resident message-passing path, // where ghost features are exchanged across ranks inside the forward pass. deepmd_compat::InputNlist make_comm_nlist(); + /** Whether this timestep requires sampling, including setup at step zero. */ + bool model_deviation_step() const; + /** Compute native-unit deviations after ghost forces have been folded. */ + void write_model_deviation( + const std::vector >& all_virial); + /** Convert and write summary statistics and optional local-order atom data. + */ + void write_model_deviation_output(const std::array& deviation, + const std::vector& std_f); private: CommBrickDeepMD* commdata_; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index f54873dbe1..20a6e11505 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -168,12 +168,76 @@ void PairDeepMDKokkos::unpack_reverse_comm_kokkos( }); } +template +bool PairDeepMDKokkos::initialize_models( + const std::vector& models) { + if (models.size() == 1) { + return PairDeepMD::initialize_models(models); + } + if (comm->nprocs != 1) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires one MPI rank; " + "use pair style deepmd for multi-rank model deviation."); + } + const int gpu_rank = get_node_rank(); + deep_pot.init(models[0], gpu_rank, get_file_content(models[0])); + std::string driver_types; + deep_pot.get_type_map(driver_types); + std::vector driver_charge_spin; + reference_models.reserve(models.size() - 1); + for (std::size_t i = 0; i < models.size(); ++i) { + deepmd_compat::DeepPot* model = &deep_pot; + if (i > 0) { + reference_models.emplace_back(new deepmd_compat::DeepPot()); + model = reference_models.back().get(); + model->init(models[i], gpu_rank, get_file_content(models[i])); + } + if (!model->uses_canonical_graph_inference() || + !model->uses_fp32_edge_vectors()) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires compatible " + "float32 compact canonical .pt2 models."); + } + if (!model->has_atomic_virial()) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires atomic virial " + "output from every model."); + } + const auto model_charge_spin = model->get_default_chg_spin(); + if (i == 0) { + driver_charge_spin = model_charge_spin; + } else if (model_charge_spin != driver_charge_spin) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires identical " + "frozen charge/spin defaults for all models."); + } + if (model->dim_fparam() != 0 || model->dim_aparam() != 0 || + model->numb_types_spin() != 0) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires models with " + "zero fparam and aparam dimensions and no spin types."); + } + std::string model_types; + model->get_type_map(model_types); + if (model->cutoff() != deep_pot.cutoff() || + model->numb_types() != deep_pot.numb_types() || + model->dim_chg_spin() != deep_pot.dim_chg_spin() || + model_types != driver_types) { + error->all( + FLERR, + "pair style deepmd/kk model deviation requires identical " + "model cutoffs and type maps, and matching charge/spin widths."); + } + } + return false; +} + template void PairDeepMDKokkos::init_style() { // Base setup and the full neighbor-list request. PairDeepMD::init_style(); - // The device edge path requires a GPU execution space and a single model. + // The driving model requires a GPU execution space. if (std::is_same::value) { error->all(FLERR, "pair style deepmd/kk runs on the GPU backend only."); } @@ -190,8 +254,17 @@ void PairDeepMDKokkos::init_style() { // set. Message-passing models additionally receive communication metadata // through compute_edges_gpu. multi_rank = (comm->nprocs > 1); - if (numb_models != 1) { - error->all(FLERR, "pair style deepmd/kk does not support model deviation."); + if (numb_models > 1 && multi_rank) { + error->all(FLERR, + "pair style deepmd/kk model deviation requires one MPI rank; " + "use pair style deepmd for multi-rank model deviation."); + } + if (numb_models > 1 && + (do_compute_fparam || do_fix_fparam || do_compute_aparam || do_ttm || + !fparam.empty() || !aparam.empty() || !charge_spin.empty())) { + error->all(FLERR, + "pair style deepmd/kk model deviation does not support runtime " + "fparam, aparam, charge/spin, compute, fix or ttm parameters."); } // A local edge graph folds ghost neighbours onto their local owner through // the atom map; without it the fold returns -1 and corrupts the graph. @@ -408,6 +481,244 @@ int PairDeepMDKokkos::build_edges_device() { return nedge; } +template +void PairDeepMDKokkos::compute_model_deviation_device() { + Kokkos::Profiling::pushRegion("deepmd/kk:model_deviation"); + const int nlocal = atom->nlocal; + const int nloc_m = compact_graph.nloc_model; + const int nnode_m = compact_graph.nnode_model; + const double natoms = static_cast(atom->natoms); + if (natoms <= 0) { + error->all(FLERR, "deepmd/kk model deviation requires at least one atom."); + } + if (d_devi_force_mean.extent(0) < static_cast(3) * nloc_m) { + d_devi_force_mean = Kokkos::View( + "deepmd/kk:devi_force_mean", static_cast(3) * nloc_m); + d_devi_force_m2 = + Kokkos::View("deepmd/kk:devi_force_m2", nloc_m); + } + if (!d_devi_summary.data()) { + d_devi_virial_sum = + Kokkos::View("deepmd/kk:devi_virial_sum"); + d_devi_virial_mean = + Kokkos::View("deepmd/kk:devi_virial_mean"); + d_devi_virial_m2 = + Kokkos::View("deepmd/kk:devi_virial_m2"); + d_devi_summary = + Kokkos::View("deepmd/kk:devi_summary"); + } + Kokkos::deep_copy(d_devi_summary, 0.0); + auto mean_f = d_devi_force_mean; + auto m2_f = d_devi_force_m2; + auto sum_v = d_devi_virial_sum; + auto mean_v = d_devi_virial_mean; + auto m2_v = d_devi_virial_m2; + auto summary = d_devi_summary; + auto out_force = d_out_force; + auto atom_virial = d_atom_virial; + + // Model 0's driving outputs are still in these buffers. Once incorporated + // in the statistics, they can be overwritten: force scatter and every + // requested LAMMPS energy/virial output have already completed above. + for (int imodel = 0; imodel < numb_models; ++imodel) { + if (imodel > 0 && nnode_m > 0) { + Kokkos::fence(); // the previous statistics must finish reading scratch + Kokkos::Profiling::pushRegion("deepmd/kk:reference_inference"); + try { + reference_models[imodel - 1]->compute_canonical_graph_gpu( + d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), + compact_graph.d_model_type.data(), compact_graph.d_source.data(), + compact_graph.d_edge_vec.data(), + compact_graph.d_destination_row_ptr.data(), + compact_graph.d_source_row_ptr.data(), + compact_graph.d_source_order.data(), nloc_m, nnode_m, + compact_graph.storage_count); + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + // The GPU API synchronizes its PyTorch stream before returning. + Kokkos::Profiling::popRegion(); + } + const double count = static_cast(imodel + 1); + Kokkos::parallel_for( + "deepmd/kk:devi_accumulate_force", + Kokkos::RangePolicy(0, nloc_m), KOKKOS_LAMBDA(const int m) { + double increment = 0.0; + for (int k = 0; k < 3; ++k) { + const double value = out_force(3 * m + k); + if (!Kokkos::isfinite(value)) { + Kokkos::atomic_exchange(&summary(6), 1.0); + } + if (imodel == 0) { + mean_f(3 * m + k) = value; + } else { + const double delta = value - mean_f(3 * m + k); + mean_f(3 * m + k) += delta / count; + increment += delta * (value - mean_f(3 * m + k)); + } + } + if (imodel == 0) { + m2_f(m) = 0.0; + } else { + m2_f(m) += increment; + } + }); + // Keep all nine reduced virial components on the device. They are needed + // on every sample, even when LAMMPS did not request a thermodynamic virial. + for (int k = 0; k < 9; ++k) { + auto component = Kokkos::subview(sum_v, k); + Kokkos::parallel_reduce( + "deepmd/kk:devi_sum_virial", + Kokkos::RangePolicy(0, nnode_m), + KOKKOS_LAMBDA(const int m, double& acc) { + const double value = atom_virial(9 * m + k); + if (!Kokkos::isfinite(value)) { + Kokkos::atomic_exchange(&summary(6), 1.0); + } + acc += value; + }, + component); + } + Kokkos::parallel_for( + "deepmd/kk:devi_accumulate_virial", + Kokkos::RangePolicy(0, 9), KOKKOS_LAMBDA(const int k) { + const double value = sum_v(k) / natoms; + if (!Kokkos::isfinite(value)) { + Kokkos::atomic_exchange(&summary(6), 1.0); + } + if (imodel == 0) { + mean_v(k) = value; + m2_v(k) = 0.0; + } else { + const double delta = value - mean_v(k); + mean_v(k) += delta / count; + m2_v(k) += delta * (value - mean_v(k)); + } + }); + } + + const double model_count = static_cast(numb_models); + const bool relative_f = out_rel == 1; + const bool relative_v = out_rel_v == 1; + const double epsilon_f = eps; + const double epsilon_v = eps_v; + const double largest = std::numeric_limits::max(); + Kokkos::parallel_for( + "deepmd/kk:devi_finish_force", Kokkos::RangePolicy(0, nloc_m), + KOKKOS_LAMBDA(const int m) { + double value = Kokkos::sqrt(m2_f(m) / model_count); + if (relative_f) { + double norm2 = 0.0; + for (int k = 0; k < 3; ++k) { + norm2 += mean_f(3 * m + k) * mean_f(3 * m + k); + } + value /= Kokkos::sqrt(norm2) + epsilon_f; + } + m2_f(m) = value; // no more models: reuse M2 for final atom deviations + if (!Kokkos::isfinite(value)) { + Kokkos::atomic_exchange(&summary(6), 1.0); + } + }); + // Include virtual (NULL-type) atoms as zeros, matching the ordinary pair's + // all-atom min/max and division by atom->natoms. + auto loc2model = compact_graph.d_loc2model; + auto fmax = Kokkos::subview(summary, 3); + auto fmin = Kokkos::subview(summary, 4); + auto fsum = Kokkos::subview(summary, 5); + Kokkos::parallel_reduce( + "deepmd/kk:devi_force_max", Kokkos::RangePolicy(0, nlocal), + KOKKOS_LAMBDA(const int i, double& acc) { + const int m = loc2model(i); + const double value = m < 0 ? 0.0 : m2_f(m); + if (value > acc) { + acc = value; + } + }, + Kokkos::Max(fmax)); + Kokkos::parallel_reduce( + "deepmd/kk:devi_force_min", Kokkos::RangePolicy(0, nlocal), + KOKKOS_LAMBDA(const int i, double& acc) { + const int m = loc2model(i); + const double value = m < 0 ? 0.0 : m2_f(m); + if (value < acc) { + acc = value; + } + }, + Kokkos::Min(fmin)); + Kokkos::parallel_reduce( + "deepmd/kk:devi_force_sum", Kokkos::RangePolicy(0, nlocal), + KOKKOS_LAMBDA(const int i, double& acc) { + const int m = loc2model(i); + acc += m < 0 ? 0.0 : m2_f(m); + }, + fsum); + Kokkos::parallel_for( + "deepmd/kk:devi_finish_summary", Kokkos::RangePolicy(0, 1), + KOKKOS_LAMBDA(const int) { + double vmax = 0.0; + double vmin = largest; + double vsq = 0.0; + for (int k = 0; k < 9; ++k) { + double value = Kokkos::sqrt(m2_v(k) / model_count); + if (relative_v) { + value /= Kokkos::abs(mean_v(k)) + epsilon_v; + } + if (!Kokkos::isfinite(value)) { + summary(6) = 1.0; + } + if (value > vmax) { + vmax = value; + } + if (value < vmin) { + vmin = value; + } + vsq += value * value; + } + summary(0) = vmax; + summary(1) = vmin; + summary(2) = Kokkos::sqrt(vsq / 9.0); + summary(5) /= natoms; + for (int k = 0; k < 6; ++k) { + if (!Kokkos::isfinite(summary(k))) { + summary(6) = 1.0; + } + } + }); + Kokkos::Profiling::pushRegion("deepmd/kk:model_deviation_summary"); + auto h_summary = Kokkos::create_mirror_view(d_devi_summary); + Kokkos::deep_copy(h_summary, d_devi_summary); + Kokkos::Profiling::popRegion(); + if (h_summary(6) != 0.0) { + error->all( + FLERR, + "Non-finite force, virial or statistic in deepmd/kk model deviation."); + } + + std::vector std_f; + if (out_each == 1) { + // This optional transfer carries final output values, never model forces. + if (d_devi_atomic.extent(0) < static_cast(nlocal)) { + d_devi_atomic = Kokkos::View( + "deepmd/kk:atomic_deviation", nlocal); + } + auto atomic_deviation = d_devi_atomic; + Kokkos::parallel_for( + "deepmd/kk:devi_atomic_output", + Kokkos::RangePolicy(0, nlocal), KOKKOS_LAMBDA(const int i) { + const int m = loc2model(i); + atomic_deviation(i) = m < 0 ? 0.0 : m2_f(m); + }); + auto h_atomic = Kokkos::create_mirror_view(d_devi_atomic); + Kokkos::deep_copy(h_atomic, d_devi_atomic); + std_f.assign(h_atomic.data(), h_atomic.data() + nlocal); + atomKK->sync(Host, TAG_MASK); + } + write_model_deviation_output({h_summary(0), h_summary(1), h_summary(2), + h_summary(3), h_summary(4), h_summary(5)}, + std_f); + Kokkos::Profiling::popRegion(); +} + template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -534,6 +845,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // them; edge-input models consume them). const double* coord_ptr = compact_graph.has_null_types ? d_coord_model.data() : x.data(); + Kokkos::Profiling::pushRegion("deepmd/kk:driver_inference"); try { if (canonical_graph) { deep_pot.compute_canonical_graph_gpu( @@ -560,6 +872,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } + Kokkos::Profiling::popRegion(); } // === Scatter the model-node forces onto their atoms === @@ -691,6 +1004,9 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { reverse_virial = false; } } + if (model_deviation_step()) { + compute_model_deviation_device(); + } } namespace LAMMPS_NS { diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index bd3689aedf..98d5cd016c 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -19,6 +19,7 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); #include #include +#include #include "compact_canonical_graph_kokkos.h" #include "kokkos_base.h" @@ -47,7 +48,11 @@ using DeepMDKokkosCommBuffer = DAT::tdual_double_1d; // returned per-atom force / energy / virial are scattered back into the // Kokkos atom arrays without any host round-trip. This removes the per-step // host coordinate marshaling and the host-device transfers of the standalone -// ``pair_style deepmd`` path. +// ``pair_style deepmd`` path. With multiple compatible canonical models, +// deviation steps reuse model 0's outputs and the same device graph for the +// reference models. Kokkos accumulates deviations on the device; only final +// statistics are copied to the host for output. This mode requires one MPI +// rank, matching cutoffs and type maps, and no runtime model parameters. // // A single rank uses the folded minimum-image node set (box thickness // > 2 * cutoff along every periodic direction); domain decomposition uses the @@ -81,8 +86,25 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // neighbor list, returning the edge count. Public because it launches // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); + /** Sample reference models after driving outputs have been consumed. + * Public because CUDA extended lambdas require a public enclosing member. + */ + void compute_model_deviation_device(); protected: + /** Load and validate the canonical ensemble without a legacy observer. */ + bool initialize_models(const std::vector& models) override; + // Each reference owns its C/C++ API object. In particular, never copy the + // C API wrapper's owning raw handle when growing this vector. + std::vector > reference_models; + Kokkos::View d_devi_force_mean; // 3 * nloc_model + Kokkos::View d_devi_force_m2; // nloc_model + Kokkos::View d_devi_atomic; // optional local output + Kokkos::View d_devi_virial_sum; + Kokkos::View d_devi_virial_mean; + Kokkos::View d_devi_virial_m2; + // max/min/rms virial, max/min/mean force, and a non-finite status flag. + Kokkos::View d_devi_summary; // Model node set and, for a compact canonical artifact, the graph itself. CompactCanonicalGraphKokkos compact_graph; Kokkos::View diff --git a/source/lmp/tests/dpa4c_model_devi_fixture.py b/source/lmp/tests/dpa4c_model_devi_fixture.py new file mode 100644 index 0000000000..8f973e762a --- /dev/null +++ b/source/lmp/tests/dpa4c_model_devi_fixture.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Subprocess-only export and inference for the official small DPA4C fixture. + +Keep PyTorch and the LAMMPS plugin in separate processes, as in the existing +``test_lammps_model_devi_pt2`` oracle. No generated archives are checked in. +""" + +import argparse +import json +import sys +import zipfile +from pathlib import ( + Path, +) + +import numpy as np + + +def export_models(directory: Path) -> None: + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + # Same standalone-helper import as source/tests/infer/gen_dpa4c_spin.py. + sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests")) + from dpa4_fixtures import ( + activate_dpa4c_condition_head, + conditioned_dpa4c_config, + jitter_zero_arrays, + ) + + directory.mkdir(parents=True, exist_ok=True) + for name, seed, state in ( + ("m0", 19, (0.0, 1.0)), + ("m1", 29, (0.0, 1.0)), + ("other_state", 19, (2.0, 3.0)), + ): + config = conditioned_dpa4c_config(state) + config["fitting_net"]["seed"] = seed + model = get_model(config).to("cpu").eval() + # Existing DPA4 fixtures activate zero residual projections so force + # and virial comparisons cannot pass on edge-independent models. + model = ( + BaseModel.deserialize( + jitter_zero_arrays(model.serialize(), np.random.default_rng(seed)) + ) + .to("cuda") + .eval() + ) + activate_dpa4c_condition_head(model.get_descriptor()) + model.get_descriptor().enable_compression(min_nbor_dist=0.5) + path = directory / f"{name}.pt2" + deserialize_to_file( + str(path), + {"model": model.serialize()}, + lower_kind="dpa4c_canonical", + do_atomic_virial=True, + ) + with zipfile.ZipFile(path) as archive: + metadata = json.loads(archive.read("model/extra/metadata.json")) + assert metadata["lower_input_kind"] == "dpa4c_canonical" + assert metadata["graph_edge_dtype"] == "float32" + assert metadata["do_atomic_virial"] is True + assert metadata["default_chg_spin"] == list(state) + assert "charge_state_constants" in metadata + + +def evaluate(request: Path, result: Path) -> None: + from deepmd.infer import ( + DeepPot, + ) + + data = json.loads(request.read_text()) + coords = np.asarray(data["coords"]) + cells = np.tile(np.eye(3).reshape(1, 9) * 13.0, (len(coords), 1)) + atom_types = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) + outputs = [] + for path in data["models"]: + model = DeepPot(path) + energy, force, virial, _, atom_virial = model.eval( + coords, cells, atom_types, atomic=True + ) + outputs.append( + { + "energy": energy.reshape(-1).tolist(), + "force": force.reshape(-1, 6, 3).tolist(), + "virial": virial.reshape(-1, 9).tolist(), + "atom_virial": atom_virial.reshape(-1, 6, 9).tolist(), + } + ) + result.write_text(json.dumps(outputs)) + + +def check_c_api(library: Path, model: Path, state: list, atomic_virial: bool) -> None: + """Exercise the public caller-owned buffer contract in isolation.""" + import ctypes + + api = ctypes.CDLL(str(library)) + api.DP_NewDeepPot.argtypes = [ctypes.c_char_p] + api.DP_NewDeepPot.restype = ctypes.c_void_p + api.DP_DeleteDeepPot.argtypes = [ctypes.c_void_p] + api.DP_DeleteDeepPot.restype = None + api.DP_DeepPotCheckOK.argtypes = [ctypes.c_void_p] + api.DP_DeepPotCheckOK.restype = ctypes.c_void_p + api.DP_DeleteChar.argtypes = [ctypes.c_void_p] + api.DP_DeleteChar.restype = None + api.DP_DeepPotGetDefaultChgSpin.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_double), + ctypes.c_int, + ] + api.DP_DeepPotGetDefaultChgSpin.restype = ctypes.c_int + api.DP_DeepPotHasAtomicVirial.argtypes = [ctypes.c_void_p] + api.DP_DeepPotHasAtomicVirial.restype = ctypes.c_bool + + def error(handle): + message = api.DP_DeepPotCheckOK(handle) + try: + return ctypes.string_at(message).decode() if message else "" + finally: + if message: + api.DP_DeleteChar(message) + + handle = api.DP_NewDeepPot(str(model).encode()) + assert handle + try: + assert not error(handle) + assert api.DP_DeepPotHasAtomicVirial(handle) is atomic_virial + assert not error(handle) + getter = api.DP_DeepPotGetDefaultChgSpin + size = getter(handle, None, 0) + assert size == len(state) + assert not error(handle) + sentinel = -12345.0 + buffer = (ctypes.c_double * (size + 1))(*([sentinel] * (size + 1))) + assert getter(handle, buffer, size) == size + assert list(buffer) == [*state, sentinel] + assert not error(handle) + # A failed query must not partially overwrite caller-owned storage. + for pointer, capacity in ( + (buffer, size - 1), + (buffer, -1), + (None, 1), + (None, -1), + ): + for index in range(size + 1): + buffer[index] = sentinel + before = list(buffer) + assert getter(handle, pointer, capacity) == -1 + assert list(buffer) == before + assert error(handle), "C API must expose the buffer error through CheckOK" + assert getter(handle, None, 0) == size + assert not error(handle), "successful length query must clear the old error" + + assert getter(handle, None, -1) == -1 + assert error(handle) + assert getter(handle, buffer, size) == size + assert list(buffer) == [*state, sentinel] + assert not error(handle), "successful buffer query must clear the old error" + + assert getter(handle, None, -1) == -1 + assert error(handle) + assert api.DP_DeepPotHasAtomicVirial(handle) is atomic_virial + assert not error(handle), ( + "successful capability query must clear the old error" + ) + finally: + api.DP_DeleteDeepPot(handle) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="action", required=True) + sub.add_parser("export").add_argument("directory", type=Path) + oracle = sub.add_parser("oracle") + oracle.add_argument("request", type=Path) + oracle.add_argument("result", type=Path) + capi = sub.add_parser("capi") + capi.add_argument("library", type=Path) + capi.add_argument("model", type=Path) + capi.add_argument("state", type=json.loads) + capi.add_argument("atomic_virial", type=json.loads) + args = parser.parse_args() + if args.action == "export": + export_models(args.directory) + elif args.action == "oracle": + evaluate(args.request, args.result) + else: + check_c_api(args.library, args.model, args.state, args.atomic_virial) diff --git a/source/lmp/tests/kokkos_model_devi_probe.cpp b/source/lmp/tests/kokkos_model_devi_probe.cpp new file mode 100644 index 0000000000..279c8e206d --- /dev/null +++ b/source/lmp/tests/kokkos_model_devi_probe.cpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Minimal Kokkos Tools observer: count existing production profiling labels. +// The callback ABI does not require linking a second Kokkos runtime. +#include +#include +#include +#include +#include + +namespace { +std::map calls; +} + +extern "C" void kokkosp_push_profile_region(const char* name) { ++calls[name]; } + +extern "C" void kokkosp_pop_profile_region() {} + +extern "C" void kokkosp_begin_parallel_for(const char* name, + const std::uint32_t, + std::uint64_t* id) { + *id = ++calls[name]; +} + +extern "C" void kokkosp_end_parallel_for(const std::uint64_t) {} + +extern "C" void kokkosp_finalize_library() { + const char* path = std::getenv("DEEPMD_KOKKOS_PROBE_OUTPUT"); + if (path == nullptr) { + return; + } + std::ofstream output(path); + for (const auto& entry : calls) { + output << entry.first << '\t' << entry.second << '\n'; + } +} diff --git a/source/lmp/tests/test_lammps_dpa4c_model_devi_kk.py b/source/lmp/tests/test_lammps_dpa4c_model_devi_kk.py new file mode 100644 index 0000000000..aab0e074df --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_model_devi_kk.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Kokkos model deviation on the official six-atom, small DPA4C fixture. + +Requires a CUDA Kokkos LAMMPS (``DEEPMD_TEST_LAMMPS`` or ``lmp`` on PATH). +``DEEPMD_TEST_PLUGIN`` optionally names its DeePMD plugin. Archives are +generated in pytest's temporary directory, or explicitly supplied through +``DEEPMD_KOKKOS_TEST_MODELS`` after running ``dpa4c_model_devi_fixture.py export``. +CPU/non-Kokkos installations skip before compiling any models. PyTorch and +LAMMPS always execute in separate subprocesses. +""" + +import json +import os +import re +import shlex +import shutil +import subprocess +import sys +import zipfile +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from write_lmp_data import ( + write_lmp_data, +) + +HERE = Path(__file__).resolve().parent +FIXTURE = HERE / "dpa4c_model_devi_fixture.py" +# Same PBC six-atom water geometry as test_lammps_dpa4_graph_pt2.py. +COORD = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +BOX = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +TYPES = np.array([1, 2, 2, 1, 2, 2]) + + +def _checked_run(argv, *, env, cwd, timeout=180, input=None): + result = subprocess.run( + argv, + input=input, + text=True, + capture_output=True, + env=env, + cwd=cwd, + timeout=timeout, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + return result + + +@pytest.fixture(scope="module") +def runtime(tmp_path_factory): + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("PyTorch support is disabled") + configured = os.environ.get("DEEPMD_TEST_LAMMPS") + executable = shutil.which(configured or "lmp") + if executable is None: + if configured: + pytest.fail(f"DEEPMD_TEST_LAMMPS is not executable: {configured}") + pytest.skip("LAMMPS executable is not available") + directory = tmp_path_factory.mktemp("dpa4c_kokkos_capabilities") + env = os.environ.copy() + env.update( + DP_CUDA_INFER="2", + OMP_NUM_THREADS="1", + DP_INTRA_OP_PARALLELISM_THREADS="1", + DP_INTER_OP_PARALLELISM_THREADS="1", + ) + # Missing hardware/optional packages is a skip; broken imports are errors. + probe = subprocess.run( + [ + sys.executable, + "-c", + "import importlib.util, sys\n" + "if importlib.util.find_spec('torch') is None: sys.exit(77)\n" + "import torch\n" + "if not torch.cuda.is_available(): sys.exit(77)\n" + "import deepmd.pt.cxx_op\n" + "from deepmd.pt_expt.utils import env\n" + "from deepmd.pt_expt.kernels.dpa4c.graph_compress import op_available\n" + "sys.exit(0 if env.DEVICE.type == 'cuda' and op_available() else 77)", + ], + capture_output=True, + text=True, + timeout=120, + env=env, + ) + if probe.returncode == 77: + pytest.skip("CUDA and the compiled DPA4C operator are required") + assert probe.returncode == 0, probe.stdout + probe.stderr + help_result = _checked_run([executable, "-h"], env=env, cwd=directory) + if "KOKKOS" not in help_result.stdout: + pytest.skip("LAMMPS was built without Kokkos") + command = [ + executable, + "-k", + "on", + "g", + "1", + "-sf", + "kk", + "-pk", + "kokkos", + "neigh", + "full", + "newton", + "off", + "gpu/aware", + "off", + "-log", + "none", + "-in", + "/dev/stdin", + ] + plugin = env.get("DEEPMD_TEST_PLUGIN") + prefix = f'plugin load "{plugin}"\n' if plugin else "" + styles = subprocess.run( + command, + input=prefix + "info styles pair\n", + text=True, + capture_output=True, + env=env, + cwd=directory, + timeout=120, + check=False, + ) + output = styles.stdout + styles.stderr + if styles.returncode != 0 and "without GPU support" in output: + pytest.skip("LAMMPS Kokkos was built without GPU support") + assert styles.returncode == 0, output + if "deepmd/kk" not in output: + pytest.skip("The loaded DeePMD plugin has no deepmd/kk pair style") + return command, env, prefix + + +@pytest.fixture(scope="module") +def models(runtime, tmp_path_factory): + _, env, _ = runtime + configured = os.environ.get("DEEPMD_KOKKOS_TEST_MODELS") + if configured: + directory = Path(configured).resolve() + else: + directory = tmp_path_factory.mktemp("dpa4c_canonical_models") + _checked_run( + [sys.executable, str(FIXTURE), "export", str(directory)], + env=env, + cwd=directory, + timeout=1200, + ) + result = {} + for name in ("m0", "m1", "other_state"): + path = directory / f"{name}.pt2" + assert path.is_file(), f"Required generated fixture is missing: {path}" + with zipfile.ZipFile(path) as archive: + metadata = json.loads(archive.read("model/extra/metadata.json")) + assert metadata["lower_input_kind"] == "dpa4c_canonical" + assert metadata["graph_edge_dtype"] == "float32" + result[name] = path + return result + + +def _archive_variant( + source, + destination, + *, + no_fold=False, + atomic_virial=None, + unknown_condition=False, + omit_atomic_output=False, +): + # Like gen_corrupt_with_comm.py: derive a temporary archive without + # rewriting or uploading a compiled model. The original stays untouched. + with ( + zipfile.ZipFile(source) as original, + zipfile.ZipFile(destination, "w") as target, + ): + for info in original.infolist(): + if no_fold and info.filename == "model/extra/charge_state.pt2": + continue + data = original.read(info.filename) + if info.filename == "model/extra/metadata.json": + metadata = json.loads(data) + if no_fold: + metadata.pop("charge_state_constants") + assert metadata["dim_chg_spin"] == 0 + assert metadata["default_chg_spin"] + if atomic_virial is not None: + metadata["do_atomic_virial"] = atomic_virial + if unknown_condition: + assert no_fold + for key in ( + "has_chg_spin_ebd", + "has_default_chg_spin", + "default_chg_spin", + "chg_spin_table_ranges", + ): + metadata.pop(key, None) + if omit_atomic_output: + assert metadata["do_atomic_virial"] is True + metadata["output_keys"].remove("atom_virial") + data = json.dumps(metadata).encode() + target.writestr(info, data) + return destination + + +def _run_lammps( + directory, + runtime, + selected, + *, + atomic=False, + out_freq=5, + rebuild=5, + nprocs=1, + probe=None, + expect_error=None, +): + directory.mkdir(parents=True, exist_ok=True) + data = directory / "water.lmp" + write_lmp_data(BOX, COORD, TYPES, data) + command, base_env, prefix = runtime + env = base_env.copy() + if probe: + env["KOKKOS_TOOLS_LIBS"] = str(probe) + env["DEEPMD_KOKKOS_PROBE_OUTPUT"] = str(directory / "profile.tsv") + options = " atomic" if atomic else "" + names = " ".join(f'"{path}"' for path in selected) + script = ( + prefix + + f"""units metal +boundary p p p +atom_style atomic +atom_modify map array sort 0 0.0 +neighbor 2.0 bin +neigh_modify every {rebuild} delay 0 check no +read_data "{data}" +mass 1 16 +mass 2 2 +pair_style deepmd/kk {names} out_file deviation.out out_freq {out_freq}{options} +pair_coeff * * +timestep 0.0005 +velocity all set 0.1 0.02 -0.03 +fix integration all nve +thermo 1 +thermo_style custom step pe +thermo_modify format float %.16e +dump trajectory all custom 1 trajectory.dump id type x y z fx fy fz +dump_modify trajectory sort id format float %.16e +run 10 +""" + ) + (directory / "input.lammps").write_text(script) + argv = list(command) + if nprocs != 1: + launcher = shutil.which("mpirun") + if launcher is None: + pytest.skip("mpirun is unavailable") + argv = [launcher, "-np", str(nprocs), *argv] + result = subprocess.run( + argv, + input=script, + text=True, + capture_output=True, + cwd=directory, + env=env, + timeout=180, + check=False, + ) + output = result.stdout + result.stderr + (directory / "lammps.out").write_text(output) + if expect_error: + assert result.returncode != 0, output + assert "ERROR" in output and expect_error in output, output + assert "Loop time" not in output, "invalid ensemble started dynamics" + return None + assert result.returncode == 0, output + assert "Loop time" in output + frames = {} + lines = (directory / "trajectory.dump").read_text().splitlines() + cursor = 0 + while cursor < len(lines): + assert lines[cursor] == "ITEM: TIMESTEP" + step = int(lines[cursor + 1]) + natoms = int(lines[cursor + 3]) + assert natoms == 6 + assert lines[cursor + 8] == "ITEM: ATOMS id type x y z fx fy fz" + table = np.array( + [ + list(map(float, line.split())) + for line in lines[cursor + 9 : cursor + 9 + natoms] + ] + ) + np.testing.assert_array_equal(table[:, 0], np.arange(1, 7)) + frames[step] = table + cursor += 9 + natoms + assert list(frames) == list(range(11)) + energies = {} + for line in output.splitlines(): + match = re.fullmatch(r"\s*(\d+)\s+([-+\d.eE]+)\s*", line) + if match: + energies[int(match[1])] = float(match[2]) + assert set(energies) == set(range(11)), output + deviation_path = directory / "deviation.out" + rows = ( + [ + line + for line in deviation_path.read_text().splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + if deviation_path.exists() + else [] + ) + deviations = ( + np.loadtxt(rows, ndmin=2) if rows else np.empty((0, 13 if atomic else 7)) + ) + return frames, energies, deviations + + +def _oracle(directory, runtime, selected, frames): + request = directory / "oracle_request.json" + result = directory / "oracle_result.json" + request.write_text( + json.dumps( + { + "models": [str(path) for path in selected], + "coords": [table[:, 2:5].tolist() for table in frames.values()], + } + ) + ) + _checked_run( + [sys.executable, str(FIXTURE), "oracle", str(request), str(result)], + env=runtime[1], + cwd=directory, + ) + outputs = json.loads(result.read_text()) + return {key: np.array([model[key] for model in outputs]) for key in outputs[0]} + + +def _assert_oracle(frames, energies, deviations, oracle, *, atomic, out_freq): + force = oracle["force"] + assert np.all(np.isfinite(force)) + assert np.max(np.abs(force)) > 1e-7, "fixture must have nonzero forces" + np.testing.assert_allclose( + [table[:, 5:8] for table in frames.values()], + force[0], + rtol=3e-4, + atol=3e-5, + ) + np.testing.assert_allclose( + list(energies.values()), oracle["energy"][0], rtol=3e-4, atol=3e-5 + ) + steps = list(range(0, 11, out_freq)) if out_freq else [] + assert deviations.shape == (len(steps), 13 if atomic else 7) + if not steps: + return + np.testing.assert_array_equal(deviations[:, 0], steps) + std_force = np.linalg.norm(np.std(force[:, steps], axis=0, ddof=0), axis=2) + std_virial = np.std(oracle["virial"][:, steps], axis=0, ddof=0) / 6 + expected = np.column_stack( + ( + std_virial.max(axis=1), + std_virial.min(axis=1), + np.sqrt(np.mean(std_virial**2, axis=1)), + std_force.max(axis=1), + std_force.min(axis=1), + std_force.mean(axis=1), + ) + ) + np.testing.assert_allclose(deviations[:, 1:7], expected, rtol=5e-4, atol=3e-5) + if atomic: + np.testing.assert_allclose(deviations[:, 7:], std_force, rtol=5e-4, atol=3e-5) + + +@pytest.mark.parametrize("atomic", [False, True]) +@pytest.mark.parametrize("order", [("m0", "m1"), ("m1", "m0"), ("m0", "m0")]) +def test_deviation_matches_python_oracle(runtime, models, tmp_path, atomic, order): + selected = [models[name] for name in order] + frames, energies, deviations = _run_lammps( + tmp_path, runtime, selected, atomic=atomic + ) + oracle = _oracle(tmp_path, runtime, selected, frames) + _assert_oracle(frames, energies, deviations, oracle, atomic=atomic, out_freq=5) + if order[0] == order[1]: + np.testing.assert_array_equal(deviations[:, 1:], 0.0) + else: + assert np.max(np.abs(oracle["force"][0] - oracle["force"][1])) > 1e-5 + assert np.max(deviations[:, 4]) > 1e-5 + + +@pytest.mark.parametrize("rebuild", [1, 5]) +def test_sampling_disabled_preserves_driver(runtime, models, tmp_path, rebuild): + selected = [models["m0"], models["m1"]] + frames, energies, deviations = _run_lammps( + tmp_path, runtime, selected, out_freq=0, rebuild=rebuild + ) + oracle = _oracle(tmp_path, runtime, selected, frames) + _assert_oracle(frames, energies, deviations, oracle, atomic=False, out_freq=0) + + +@pytest.fixture(scope="module") +def profiling_library(runtime, tmp_path_factory): + directory = tmp_path_factory.mktemp("dpa4c_kokkos_tools") + compiler = shlex.split(os.environ.get("CXX", "c++")) + if not compiler or shutil.which(compiler[0]) is None: + pytest.skip("C++ compiler is required for the Kokkos Tools observer") + library = directory / "model_devi_probe.so" + _checked_run( + [ + *compiler, + "-std=c++11", + "-shared", + "-fPIC", + str(HERE / "kokkos_model_devi_probe.cpp"), + "-o", + str(library), + ], + env=runtime[1], + cwd=directory, + ) + return library + + +@pytest.mark.parametrize("out_freq", [0, 5]) +def test_reference_reuses_driver_graph( + runtime, models, profiling_library, tmp_path, out_freq +): + _run_lammps( + tmp_path, + runtime, + [models["m0"], models["m1"]], + out_freq=out_freq, + probe=profiling_library, + ) + counts = {} + for line in (tmp_path / "profile.tsv").read_text().splitlines(): + name, count = line.split("\t") + counts[name] = int(count) + assert counts["deepmd/kk:driver_inference"] == 11 + assert counts["compact_canonical:count"] == 11 + assert counts.get("deepmd/kk:reference_inference", 0) == (3 if out_freq else 0) + assert counts.get("deepmd/kk:model_deviation", 0) == (3 if out_freq else 0) + + +def test_multiple_mpi_ranks_rejected(runtime, models, tmp_path): + _run_lammps( + tmp_path, + runtime, + [models["m0"], models["m1"]], + nprocs=2, + expect_error="requires one MPI rank", + ) + + +@pytest.mark.parametrize("no_fold", [False, True]) +def test_different_frozen_conditions_rejected(runtime, models, tmp_path, no_fold): + selected = [models["m0"], models["other_state"]] + if no_fold: + selected = [ + _archive_variant(path, tmp_path / f"legacy_{i}.pt2", no_fold=True) + for i, path in enumerate(selected) + ] + _run_lammps(tmp_path, runtime, selected, expect_error="charge/spin") + + +def test_equal_legacy_frozen_conditions(runtime, models, tmp_path): + selected = [ + _archive_variant(models[name], tmp_path / f"legacy_{name}.pt2", no_fold=True) + for name in ("m0", "m1") + ] + frames, energies, deviations = _run_lammps(tmp_path, runtime, selected, atomic=True) + # The modern archives contain exactly the same frozen inference constants. + oracle = _oracle(tmp_path, runtime, [models["m0"], models["m1"]], frames) + _assert_oracle(frames, energies, deviations, oracle, atomic=True, out_freq=5) + + +@pytest.mark.parametrize("unsupported_index", [0, 1]) +def test_unknown_frozen_condition_rejected( + runtime, models, tmp_path, unsupported_index +): + selected = [models["m0"], models["m1"]] + selected[unsupported_index] = _archive_variant( + selected[unsupported_index], + tmp_path / "unknown_frozen_condition.pt2", + no_fold=True, + unknown_condition=True, + ) + _run_lammps( + tmp_path, + runtime, + selected, + expect_error="default charge/spin state cannot be established", + ) + + +@pytest.mark.parametrize("unsupported_index", [0, 1]) +def test_atomic_virial_required_for_every_model( + runtime, models, tmp_path, unsupported_index +): + selected = [models["m0"], models["m1"]] + selected[unsupported_index] = _archive_variant( + selected[unsupported_index], + tmp_path / "no_atomic_virial.pt2", + atomic_virial=False, + ) + _run_lammps(tmp_path, runtime, selected, expect_error="atomic virial") + + +@pytest.mark.parametrize("unsupported_index", [0, 1]) +def test_atomic_virial_output_key_required( + runtime, models, tmp_path, unsupported_index +): + selected = [models["m0"], models["m1"]] + selected[unsupported_index] = _archive_variant( + selected[unsupported_index], + tmp_path / "missing_atomic_virial_output.pt2", + omit_atomic_output=True, + ) + _run_lammps(tmp_path, runtime, selected, expect_error="atomic virial") + + +@pytest.mark.parametrize("no_fold", [False, True]) +@pytest.mark.parametrize("atomic_virial", [False, True]) +def test_c_api_metadata_buffer_contract( + runtime, models, tmp_path, no_fold, atomic_virial +): + configured = os.environ.get("DEEPMD_TEST_C_API_LIBRARY") + if not configured: + pytest.skip("Set DEEPMD_TEST_C_API_LIBRARY to exercise the C API library") + library = Path(configured).resolve() + assert library.is_file(), f"C API library is missing: {library}" + model = _archive_variant( + models["other_state"], + tmp_path / "c_api_model.pt2", + no_fold=no_fold, + atomic_virial=atomic_virial, + ) + _checked_run( + [ + sys.executable, + str(FIXTURE), + "capi", + str(library), + str(model), + "[2.0, 3.0]", + json.dumps(atomic_virial), + ], + env=runtime[1], + cwd=tmp_path, + ) diff --git a/source/tests/dpa4_fixtures.py b/source/tests/dpa4_fixtures.py index a99caf2589..a305cd0906 100644 --- a/source/tests/dpa4_fixtures.py +++ b/source/tests/dpa4_fixtures.py @@ -11,6 +11,71 @@ import numpy as np +def dpa4c_config() -> dict: + """Small DPA4C model shared by graph-lower and LAMMPS tests.""" + return { + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 16, + "lmax": 4, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + }, + } + + +def compressed_dpa4c_config(channels: int = 8) -> dict: + """The existing fp32 compact-canonical export fixture.""" + config = dpa4c_config() + config["descriptor"].update( + channels=channels, lmax=2, n_radial=8, precision="float32" + ) + # The fused fitting operator has no per-layer timestep. + config["fitting_net"].update( + neuron=[32, 32], + activation_function="silu", + precision="float32", + resnet_dt=False, + ) + return config + + +def conditioned_dpa4c_config(default_chg_spin=(2.0, 3.0)) -> dict: + """Canonical fixture with the condition used by the baked-state tests.""" + config = compressed_dpa4c_config() + config["descriptor"].update( + add_chg_spin_ebd=True, default_chg_spin=list(default_chg_spin) + ) + return config + + +def activate_dpa4c_condition_head(descriptor) -> None: + """Use the existing deterministic, nonzero charge-conditioning head.""" + import torch + + head = descriptor.charge_spin_embedding.network.layers[-1] + generator = torch.Generator(device=head.w.device).manual_seed(11) + with torch.no_grad(): + head.w.copy_( + torch.randn( + head.w.shape, + dtype=head.w.dtype, + device=head.w.device, + generator=generator, + ) + * 0.5 + ) + + def jitter_zero_arrays(node, rng: np.random.Generator): """Return a copy of a serialized tree with every zero float array jittered. diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index 5734b011fa..a5a74f4a65 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -34,6 +34,10 @@ op_available, ) +from ...dpa4_fixtures import ( + activate_dpa4c_condition_head, +) + _GPU = pytest.mark.skipif( not torch.cuda.is_available() or not op_available(), reason="CUDA and the compiled DPA4C operator are required", @@ -439,18 +443,7 @@ def _build_charge_descriptor( .cuda() .eval() ) - head = descriptor.charge_spin_embedding.network.layers[-1] - generator = torch.Generator(device="cuda").manual_seed(11) - with torch.no_grad(): - head.w.copy_( - torch.randn( - head.w.shape, - dtype=torch.float32, - device="cuda", - generator=generator, - ) - * 0.5 - ) + activate_dpa4c_condition_head(descriptor) return descriptor diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index d8f9ed31eb..1d40d05e82 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -20,6 +20,12 @@ build_synthetic_graph_inputs, ) +from ...dpa4_fixtures import compressed_dpa4c_config as _compressed_config +from ...dpa4_fixtures import ( + conditioned_dpa4c_config, +) +from ...dpa4_fixtures import dpa4c_config as _config + #: The compact canonical ABI and the fused spin backward are CUDA-only routes, #: so they follow the configured backend device rather than the mere presence #: of CUDA hardware: a run pinned to the CPU takes the generic graph lower. @@ -29,45 +35,6 @@ ) -def _config() -> dict: - return { - "type_map": ["A", "B"], - "descriptor": { - "type": "dpa4c", - "rcut": 3.0, - "channels": 16, - "lmax": 4, - "n_radial": 4, - "precision": "float64", - "seed": 17, - }, - "fitting_net": { - "type": "ener", - "neuron": [16, 16], - "precision": "float64", - "seed": 19, - }, - } - - -def _compressed_config(channels: int = 8) -> dict: - config = _config() - descriptor = config["descriptor"] - descriptor["channels"] = channels - descriptor["lmax"] = 2 - descriptor["n_radial"] = 8 - descriptor["precision"] = "float32" - fitting = config["fitting_net"] - fitting["neuron"] = [32, 32] - fitting["activation_function"] = "silu" - fitting["precision"] = "float32" - # The fused fitting operator has no per-layer timestep, and the shipped - # DPA4C grades do not use one either, so the compact canonical path is - # only reachable without it. - fitting["resnet_dt"] = False - return config - - def _run_graph( model: torch.nn.Module, dtype: torch.dtype = torch.float64, @@ -520,9 +487,7 @@ def test_an_uncompressed_export_keeps_the_charge_state_as_a_runtime_input() -> N conditioning slot with a dynamic frame axis, and only the fold of the compact canonical path removes it. """ - config = _compressed_config() - config["descriptor"]["add_chg_spin_ebd"] = True - config["descriptor"]["default_chg_spin"] = [2.0, 3.0] + config = conditioned_dpa4c_config() model = get_model(config).to("cpu").eval() exported, metadata, _model_json, _output_keys = _trace_and_export( {"model": model.serialize()}, @@ -549,9 +514,7 @@ def test_a_baked_charge_state_reaches_the_compact_canonical_lower() -> None: charge state into the frozen tables and the snapshot then reports a zero runtime condition width. """ - config = _compressed_config() - config["descriptor"]["add_chg_spin_ebd"] = True - config["descriptor"]["default_chg_spin"] = [2.0, 3.0] + config = conditioned_dpa4c_config() model = get_model(config).to("cpu").eval() assert model.get_dim_chg_spin() == 2 assert _resolve_lower_kind("model.pt2", {"model": model.serialize()}, "auto") == (