diff --git a/CHANGELOG.md b/CHANGELOG.md index 2078217..d4f635f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,43 @@ Notable changes of grace-runtime, grace-runtime-cuda12 and grace-runtime-rocm6 (they share one version). Each release has a section "Changed defaults" when a default behaves differently. +## [0.1.2](https://github.com/ICAMS/grace-runtime/releases/tag/0.1.2) (2026-11-01) + +The early GRACE-1L/2L models, and fast GPU kernels for more GRACE-1L/2L fits. + +### Added + +- **Early GRACE-1L/2L models:** GRACE-1L-OAM, GRACE-1L-OMAT, GRACE-2L-OAM and GRACE-2L-OMAT (and their older aliases) + are in the registry and run on the CPU, NVIDIA and AMD GPUs, in Python and in LAMMPS, without an uncertainty head. + Own fits of the early GRACE-1L/2L presets run as well: tensorpotential exports them as `kokkos.npz` format 2. +- **More GRACE-1L/2L fits on the fast GPU kernels.** A fit whose layers are narrower than compiled GPU kernels runs + zero-padded to them, with the same results and a warning that names the widths (not models with an uncertainty + head). GRACE-2L has kernels for the widths of the current registry models, of the early ones, and of fits with 48 + layer-2 radial channels; GRACE-1L fits run with up to 64 radial channels. The models guide lists the widths. +- `GRACE_GENERIC_KERNELS=1`: a GRACE-2L fit wider than all compiled GPU kernels runs on generic GPU kernels, several + times slower, with a warning. Without the variable such a fit stops with an error that names its widths. + +### LAMMPS (`grace-runtime` branch of the LAMMPS fork) + +- `grace/runtime` and `grace/runtime/kk` run the early GRACE-1L/2L models and own fits of their presets. The Kokkos + style pads narrower fits like Python does, and with `GRACE_GENERIC_KERNELS=1` runs wider GRACE-2L fits on the + generic GPU kernels. Warnings appear as LAMMPS `WARNING` lines. + +### Documentation + +- Benchmarks: H200, MI250X, L40S and CPU charts; the front page shows the execution speed. +- Models: which fits run on the fast GPU kernels, padded or not, and `GRACE_GENERIC_KERNELS`. + +### Changed defaults + +- A GRACE-1L/2L model without GPU kernels of its own widths stopped with an error on a GPU. A narrower one now runs + zero-padded with a warning; a wider GRACE-2L one still stops, and the error names `GRACE_GENERIC_KERNELS=1`. +- `device="auto"`: such a wider GRACE-2L model raises instead of running on the CPU. +- `kokkos.npz` format version 2 for the early GRACE-1L/2L layouts. grace-runtime 0.1.1 and older refuse these files + and ask for an update. Files of current models stay format 1. +- The GPU modules of `grace-runtime-cuda12` / `grace-runtime-rocm6` 0.1.2 have interface version 7; grace-runtime + 0.1.2 refuses older modules. The packages pin each other, so a normal `pip install` matches them. + ## [0.1.1](https://github.com/ICAMS/grace-runtime/releases/tag/0.1.1) (2026-10-29) Bug fixes and documentation. diff --git a/docs/lammps.md b/docs/lammps.md index e56ae46..adc01c2 100644 --- a/docs/lammps.md +++ b/docs/lammps.md @@ -30,7 +30,7 @@ cd lammps ## 2. Build The pair styles are part of the `ML-PACE` package. When CMake configures the build, it downloads grace-runtime -0.1.1 from GitHub, so the configure step needs network access. It also downloads the ML-PACE library. Pick one of the +0.1.2 from GitHub, so the configure step needs network access. It also downloads the ML-PACE library. Pick one of the three recipes; run each from the `lammps` folder. ### CPU @@ -131,10 +131,10 @@ Other AMD GPUs ([GPU backends](backends.md#amd-rocm)): ### Without network access -Download grace-runtime 0.1.1 on a machine with network access, copy it over and pass its path: +Download grace-runtime 0.1.2 on a machine with network access, copy it over and pass its path: ```bash -git clone --depth 1 -b 0.1.1 https://github.com/ICAMS/grace-runtime.git +git clone --depth 1 -b 0.1.2 https://github.com/ICAMS/grace-runtime.git cmake -S cmake -B build-cpu ... -D LOCAL_GRACE_RUNTIME=$PWD/grace-runtime ``` diff --git a/docs/models.md b/docs/models.md index f76e1c0..657ec2c 100644 --- a/docs/models.md +++ b/docs/models.md @@ -12,9 +12,9 @@ grace-runtime models download GRACE-2L-OMAT-medium-ft-E ``` OMAT models use a 6 Å cutoff; SMAX variants use element-dependent cutoffs of 5–7.5 Å. -Older registry aliases are accepted. Legacy GRACE-FS and early OAM/OMAT models without a Kokkos export require -tensorpotential's `TPCalculator` instead. PACE (ACE) models fitted with pacemaker run from their `.yaml`: see -[PACE models](pace.md). +Older registry aliases are accepted. The early models GRACE-1L-OAM, GRACE-1L-OMAT, GRACE-2L-OAM and GRACE-2L-OMAT +run as well, without an uncertainty head. Legacy GRACE-FS models require tensorpotential's `TPCalculator`. PACE (ACE) models fitted +with pacemaker run from their `.yaml`: see [PACE models](pace.md). To use your own model, export it with tensorpotential and pass the resulting file: @@ -22,6 +22,19 @@ To use your own model, export it with tensorpotential and pass the resulting fil grace_utils -p model.yaml -c CHECKPOINT export_kokkos -o model.npz ``` +Fits of the early GRACE-1L/2L presets export the same way. Their files are format 2; an older grace-runtime refuses +them and asks for an update. + +On a GPU, GRACE-2L fits up to the widths of the registry models run on the fast kernels, narrower layers zero-padded +with the same results and a warning (lmax at most 4; not models with an uncertainty head): a fit of the current presets +with at most 42 layer-1 and 32 layer-2 radial channels, 64 layer-2 product functions and `n_mlp_dens` at most 16, and +any fit with at most 32 layer-1 and 48 layer-2 radial channels, 48 layer-2 product functions and `n_mlp_dens` at most +12. GRACE-1L fits with up to 64 radial channels and FC outputs run zero-padded the same way. Wider GRACE-2L fits stop +with an error that names their widths, also with `device="auto"`; with `GRACE_GENERIC_KERNELS=1` set in the environment +they run on generic GPU kernels, several times slower, with a warning. Please open an +[issue](https://github.com/ICAMS/grace-runtime/issues) with the widths the message names so that fast kernels can be +added. + After exporting `model.npz`, load it with: ```python diff --git a/libgrace/include/grace/fused_1l.h b/libgrace/include/grace/fused_1l.h index 0e2c4e3..cde5ee1 100644 --- a/libgrace/include/grace/fused_1l.h +++ b/libgrace/include/grace/fused_1l.h @@ -77,6 +77,8 @@ class Grace1LFused { std::vector e_norm; std::vector e_nin, e_nout; int e_act = 0; + bool e_ln = false; // legacy models: RMS layer norm of the MLP input (ln_fwd) + bool aaa_a = false; // legacy models: AAA = AA1 x A (else AA1 x A1) double out_scale = 1.0; std::vector order, bstart, bcount; diff --git a/libgrace/include/grace/fused_2l.h b/libgrace/include/grace/fused_2l.h index b5e005b..5c29297 100644 --- a/libgrace/include/grace/fused_2l.h +++ b/libgrace/include/grace/fused_2l.h @@ -117,10 +117,13 @@ class Grace2LFused { // RMSNorms + energy MLP (double) std::vector lnI_scale, ln0_scale; + bool rms = true; // the RMSNorms (off: legacy models, unit scales) + int nrms1 = 0, nrms0 = 0; // the RMS means of rho[1:] / I2 divide by these (padded model: the unpadded widths) std::vector> eW; std::vector e_norm; std::vector e_nin, e_nout; int e_act = 1; + bool e_ln = false; // legacy models: RMS layer norm of the MLP input (ln_fwd) double out_scale = 1.0; // per-atom globals diff --git a/libgrace/include/grace/fused_cpu.h b/libgrace/include/grace/fused_cpu.h index fd43a09..0a4ed49 100644 --- a/libgrace/include/grace/fused_cpu.h +++ b/libgrace/include/grace/fused_cpu.h @@ -243,57 +243,6 @@ inline float dot3(const float *__restrict x, const float *__restrict y, const fl inline void vzero(float *y, const size_t n) { std::memset(y, 0, n * sizeof(float)); } -// fp64: Chebyshev T_1..T_nrb of (2x-1) times fcut_p(x), and d/dx -inline void radial_basis_x(double x, int p, int nrb, std::vector &g, std::vector &dg) -{ - const double xp = std::pow(x, p), xp1 = xp * x, xp2 = xp1 * x; - const double fc = 1.0 - 0.5 * (p + 1) * (p + 2) * xp + p * (p + 2) * xp1 - 0.5 * p * (p + 1) * xp2; - const double dfc = 0.5 * p * (p + 1) * (p + 2) * (-std::pow(x, p - 1) + 2.0 * xp - xp1); - const double t = 2.0 * x - 1.0; - double Tp = 1.0, Tc = t, dTp = 0.0, dTc = 2.0; - g.resize(nrb); - dg.resize(nrb); - for (int k = 0; k < nrb; k++) { - g[k] = Tc * fc; - dg[k] = dTc * fc + Tc * dfc; - const double Tn = 2.0 * t * Tc - Tp, dTn = 2.0 * (2.0 * Tc + t * dTc) - dTp; - Tp = Tc; dTp = dTc; Tc = Tn; dTc = dTn; - } -} - -// fp64 forward-mode radial MLP (silu hidden layers, linear output) -template -void mlp_x(const Layers &layers, std::vector h, std::vector dh, - std::vector &out, std::vector &dout) -{ - const int n = layers.size(); - for (int il = 0; il < n; il++) { - const auto &L = layers[il]; - std::vector s(L.n_out, 0.0), ds(L.n_out, 0.0); - for (int k = 0; k < L.n_in; k++) - for (int j = 0; j < L.n_out; j++) { - const double w = L.W[(size_t) k * L.n_out + j]; - s[j] += w * h[k]; - ds[j] += w * dh[k]; - } - for (int j = 0; j < L.n_out; j++) { - s[j] *= L.norm; - ds[j] *= L.norm; - if (il < n - 1) { - const double sg = 1.0 / (1.0 + std::exp(-s[j])); - ds[j] = sg * (1.0 + s[j] * (1.0 - sg)) * ds[j]; - s[j] = s[j] * sg; - } - } - h.swap(s); - dh.swap(ds); - } - out.swap(h); - dout.swap(dh); -} - - - // 16-float vectors (one AVX-512 register, two halves on AVX2: grace_cpu_simd.h). Unaligned access. typedef simd::f16 v16; inline v16 ld(const float *p) { return simd::load(p); } @@ -712,6 +661,25 @@ inline void red_bwd(const Red &R, int ty, const float *gOut, float *const *gX) default: break; \ } +// legacy GRACE-1L/2L: the RMS layer norm z = x q of the energy-MLP input x[0..n) in place (its scale folded into the +// first layer), q = 1 / sqrt(mean x^2 + 1e-16) as tensorpotential's scalar_rms_ln; x0 keeps the input. Returns q. +inline double ln_fwd(double *x, double *x0, int n) +{ + double ms = 0.0; + for (int k = 0; k < n; k++) ms += (x0[k] = x[k]) * x[k]; + const double q = 1.0 / std::sqrt(ms / n + 1e-16); + for (int k = 0; k < n; k++) x[k] *= q; + return q; +} + +// its reverse in place: g = dE/dz -> dE/dx +inline void ln_bwd(double *g, const double *x0, int n, double q) +{ + double t = 0.0; + for (int k = 0; k < n; k++) t += g[k] * x0[k]; + for (int k = 0; k < n; k++) g[k] = q * (g[k] - x0[k] * t * q * q / n); +} + // ---- bond geometry: neighbours within the cutoff, unit vectors, real spherical harmonics Y[lm] and the // gradient G[lm][3] of their polynomial extension (dual numbers). Shared by the CPU GRACE-1L and 2L drivers. diff --git a/libgrace/include/grace/kk/fused_common.h b/libgrace/include/grace/kk/fused_common.h index 5cab8cb..005ed1a 100644 --- a/libgrace/include/grace/kk/fused_common.h +++ b/libgrace/include/grace/kk/fused_common.h @@ -19,6 +19,11 @@ namespace grace { namespace kk { namespace GraceFusedKK { +// the last sentence of the warnings for a model without GPU kernels of its own widths +constexpr const char *GRACE_KK_ISSUE = + "Please open an issue at https://github.com/ICAMS/grace-runtime/issues with these numbers so that fast kernels " + "can be added"; + constexpr int GRACE_FUSED_MAX_LMAX = 4; struct alignas(16) GF4 { float x, y, z, w; }; @@ -185,50 +190,6 @@ void spline_vd(const View3 &c, const int N, const float x, const int o, float &v dv = (a1 + t * (2.0f * a2 + t * 3.0f * a3)) * (float) N; } -// Host fp64: Chebyshev T_1..T_nrb of (2x-1) times fcut_p(x), and d/dx (x = r / rcut_ij). -inline void radial_basis_x(double x, int p, int nrb, std::vector &g, std::vector &dg) -{ - const double xp = std::pow(x, p), xp1 = xp * x, xp2 = xp1 * x; - const double fc = 1.0 - 0.5 * (p + 1) * (p + 2) * xp + p * (p + 2) * xp1 - 0.5 * p * (p + 1) * xp2; - const double dfc = 0.5 * p * (p + 1) * (p + 2) * (-std::pow(x, p - 1) + 2.0 * xp - xp1); - const double t = 2.0 * x - 1.0; - double Tp = 1.0, Tc = t, dTp = 0.0, dTc = 2.0; // d/dx, chain factor 2 - g.resize(nrb); dg.resize(nrb); - for (int k = 0; k < nrb; k++) { - g[k] = Tc * fc; - dg[k] = dTc * fc + Tc * dfc; - const double Tn = 2.0 * t * Tc - Tp, dTn = 2.0 * (2.0 * Tc + t * dTc) - dTp; - Tp = Tc; dTp = dTc; Tc = Tn; dTc = dTn; - } -} - -// Host fp64 forward-mode radial MLP (silu hidden layers, linear output), as TagComputeMLPRadial_R. -template -void mlp_x(const std::vector &layers, std::vector h, - std::vector dh, std::vector &out, std::vector &dout) -{ - const int n = layers.size(); - for (int il = 0; il < n; il++) { - const auto &L = layers[il]; - std::vector s(L.n_out, 0.0), ds(L.n_out, 0.0); - for (int k = 0; k < L.n_in; k++) - for (int j = 0; j < L.n_out; j++) { - const double w = L.W[(size_t) k * L.n_out + j]; - s[j] += w * h[k]; ds[j] += w * dh[k]; - } - for (int j = 0; j < L.n_out; j++) { - s[j] *= L.norm; ds[j] *= L.norm; - if (il < n - 1) { - const double sg = 1.0 / (1.0 + std::exp(-s[j])); - ds[j] = sg * (1.0 + s[j] * (1.0 - sg)) * ds[j]; - s[j] = s[j] * sg; - } - } - h.swap(s); dh.swap(ds); - } - out.swap(h); dout.swap(dh); -} - } // namespace GraceFusedKK } // namespace kk } // namespace grace diff --git a/libgrace/include/grace/kk/grace_1l_fused_kk.h b/libgrace/include/grace/kk/grace_1l_fused_kk.h index 50a9068..deb0319 100644 --- a/libgrace/include/grace/kk/grace_1l_fused_kk.h +++ b/libgrace/include/grace/kk/grace_1l_fused_kk.h @@ -28,6 +28,15 @@ class Grace1LFusedKK : public Grace1LKK { // checks the model against the compiled kernels, builds the dense metadata and the radial tables void setup() override; + // a model without UQ whose radial or FC width has no compiled kernels runs zero-padded (GRACE1LModel::padded: the + // same results), with a warning + void set_model(std::shared_ptr model, const std::vector &type2mu) override; + // the widths {radial channels, FC outputs} set_model pads a model to; {0, 0}: none + static std::array pad_widths(const GRACE1LModel &m); + // the model zero-padded to its pad_widths; the model itself without them + static std::shared_ptr padded_model(std::shared_ptr model); + // the warning for a model that runs zero-padded; empty otherwise + static std::string slow_notice(const GRACE1LModel &m); typename Base::RunOutput run(const typename Base::RunInput &in) override; void release_buffers() override { diff --git a/libgrace/include/grace/kk/grace_1l_fused_kk_impl.h b/libgrace/include/grace/kk/grace_1l_fused_kk_impl.h index acec0b5..0c4e808 100644 --- a/libgrace/include/grace/kk/grace_1l_fused_kk_impl.h +++ b/libgrace/include/grace/kk/grace_1l_fused_kk_impl.h @@ -6,6 +6,7 @@ #include "grace/kk/grace_1l_kk_impl.h" #include "grace/kk/fused_common.h" #include "grace/kk/fused_dense.h" +#include "grace/radial_x.h" #include "grace/kk/fused_hip_bonds.h" #include "grace/kk/g1l_shape_large.h" @@ -28,8 +29,10 @@ constexpr int G1LF_MAX_ACC = 16; // per-thread accumulators: nradmax * nlm <= constexpr int G1LF_NO = 64; // FC widths compiled for: 64 (GRACE-1L large) and 32 (medium) constexpr int G1LF_NO_M = 32; constexpr int G1LF_NIN = 64; // max channel width of the rho sources -constexpr int G1LF_NIR = 32; // max right-input (A) width of the FCs +constexpr int G1LF_NIR = 32; // right-input (A) width of the generic reverse FC (32 and 64: fc_bwd_rt) +constexpr int G1LF_NRMAX = 64; // max radial width (zero-padded to 32 or 64 when the kernels need it) constexpr int G1LF_CC = 8; // rho connections per thread (forward) +constexpr int G1LF_RQ = 5; // rho output quads: rho_n_out <= 20 (17; legacy models 13), weights padded to 20 } // namespace g1lf_detail using namespace g1lf_detail; @@ -39,6 +42,54 @@ using namespace g1lf_detail; /* ---------------------------------------------------------------------- */ +// FC width no / radial width nr (nlm harmonics) runs on the compiled kernels +inline bool g1lf_fc_ok(int no) { return no == G1LF_NO || no == G1LF_NO_M; } +inline bool g1lf_radial_ok(int nr, int nlm) +{ + return nr <= G1LF_NRMAX && G1LF_TEAM % nr == 0 && nr * nlm <= G1LF_TEAM * G1LF_MAX_ACC; +} + +template +std::array Grace1LFusedKK::pad_widths(const GRACE1LModel &m) +{ + const int nlm = (m.lmax + 1) * (m.lmax + 1), fc = m.fc_A1.n_out; + if (m.has_uq || m.lmax > G1LF_MAX_LMAX || m.rho_n_out > 4 * G1LF_RQ || + (g1lf_radial_ok(m.nradmax, nlm) && g1lf_fc_ok(fc))) + return {0, 0}; + // radial and FC widths padded to the compiled FC widths 32 / 64 + int nr = g1lf_radial_ok(m.nradmax, nlm) ? m.nradmax : m.nradmax <= G1LF_NO_M ? G1LF_NO_M : G1LF_NO; + int no = fc <= G1LF_NO_M ? G1LF_NO_M : G1LF_NO; + if (m.AAA_right_A) nr = no = std::max(nr, no); // legacy AAA = AA1 x A: one channel count + if (m.nradmax > nr || fc > no || !g1lf_radial_ok(nr, nlm)) return {0, 0}; + return {nr, no}; +} + +template +std::shared_ptr Grace1LFusedKK::padded_model(std::shared_ptr model) +{ + const std::array w = model ? pad_widths(*model) : std::array{0, 0}; + return w[0] ? std::make_shared(model->padded(w[0], w[1])) : model; +} + +template +std::string Grace1LFusedKK::slow_notice(const GRACE1LModel &m) +{ + const std::array w = pad_widths(m); + if (!w[0]) return ""; + return fmt("this GRACE-1L model (lmax {}, radial channels {}, FC width {}) has no GPU kernels of its own widths: it " + "runs zero-padded to {} / {} (the same results, at the cost of the wider model). {}", m.lmax, m.nradmax, + m.fc_A1.n_out, w[0], w[1], GRACE_KK_ISSUE); +} + +template +void Grace1LFusedKK::set_model(std::shared_ptr model, + const std::vector &type2mu) +{ + auto p = padded_model(model); + if (p != model) this->log.warn("GRACE-1L/KK: " + slow_notice(*model)); + Base::set_model(std::move(p), type2mu); +} + template void Grace1LFusedKK::setup() { @@ -55,22 +106,20 @@ void Grace1LFusedKK::setup() const int nlm = (this->lmax + 1) * (this->lmax + 1); if (this->lmax > G1LF_MAX_LMAX) GRACE_KK_ERROR("GRACE-1L/KK: lmax {} > {} not supported", this->lmax, G1LF_MAX_LMAX); - if (this->nradmax * nlm > G1LF_TEAM * G1LF_MAX_ACC || this->nradmax > G1LF_TEAM || G1LF_TEAM % this->nradmax) + if (!g1lf_radial_ok(this->nradmax, nlm)) GRACE_KK_ERROR("GRACE-1L/KK: radial width {} outside the compiled bond kernels", this->nradmax); - if (this->fc_n_out != G1LF_NO && this->fc_n_out != G1LF_NO_M) + if (!g1lf_fc_ok(this->fc_n_out)) GRACE_KK_ERROR("GRACE-1L/KK: FC width {} not compiled (must be {} or {})", this->fc_n_out, G1LF_NO, G1LF_NO_M); - if (this->rho_n_out != 17) GRACE_KK_ERROR("GRACE-1L/KK: rho width must be 17"); - if (this->nradmax > G1LF_NIN || this->nradmax > G1LF_NIR) - GRACE_KK_ERROR("GRACE-1L/KK: nradmax {} > {}", this->nradmax, G1LF_NIR); + if (this->rho_n_out > 4 * G1LF_RQ) GRACE_KK_ERROR("GRACE-1L/KK: rho width {} > {}", this->rho_n_out, 4 * G1LF_RQ); auto live = [](const GRACE1LModel::CGProduct &p, bool left, bool right) { std::vector v; if (left) v.insert(v.end(), p.left_ind.begin(), p.left_ind.end()); if (right) v.insert(v.end(), p.right_ind.begin(), p.right_ind.end()); return v; }; - // A1 feeds AA (both operands) and AAA (right operand) - std::vector liveA1 = live(m.prod_AA, true, true), r = live(m.prod_AAA, false, true); + // A1 feeds AA (both operands) and AAA (right operand; legacy models: A) + std::vector liveA1 = live(m.prod_AA, true, true), r = live(m.prod_AAA, false, !m.AAA_right_A); liveA1.insert(liveA1.end(), r.begin(), r.end()); build_fc_meta(fcm_A1, m.fc_A1, liveA1); build_fc_meta(fcm_AA1, m.fc_AA1, live(m.prod_AAA, true, false)); @@ -124,7 +173,7 @@ void Grace1LFusedKK::build_spline_tables() std::vector g, dg, out, dout; for (int k = 0; k <= N; k++) { radial_basis_x(k * h, p, nrb, g, dg); - mlp_x(layers, g, dg, out, dout); + mlp_x(layers, g, dg, out, dout, m.radial_activation); for (int o = 0; o < nout; o++) { fv[(size_t) k * nout + o] = out[o]; dv[(size_t) k * nout + o] = h * dout[o]; } } d_splR = V3("g1lf:R", N, nout, 4); @@ -590,7 +639,7 @@ void Grace1LFusedKK::build_red_meta() const GRACE1LModel &m = *this->grace_model; const GRACE1LModel::ReduceWeights *rws[4] = {&m.reduce_A, &m.reduce_AA, &m.reduce_AAA, &m.reduce_AAAA}; const int nfi[4] = {(this->lmax + 1) * (this->lmax + 1), this->n_funcs_AA, this->n_funcs_AAA, this->n_funcs_AAAA}; - const int ntypes = ((int) this->h_type2model.size() - 1), no = this->rho_n_out, nop = (no + 3) / 4 * 4; + const int ntypes = ((int) this->h_type2model.size() - 1), no = this->rho_n_out, nop = 4 * G1LF_RQ; rm.nx = 4; int gtot = 0; std::vector> conns; // (x, fi, tile), sorted by x @@ -808,8 +857,9 @@ void Grace1LFusedKK::stage_dense_fwd() else prod_fwd(pm_AA, this->d_A1, this->d_A1, this->d_AA); if (this->fc_n_out == G1LF_NO) fc_fwd4(fcm_AA1, this->d_AA, this->d_A, this->d_AA1); else fc_fwd4(fcm_AA1, this->d_AA, this->d_A, this->d_AA1); - if (shape_id) G1LF_SHAPE_FWD(AAA, this->d_AA1, this->d_A1, this->d_AAA) - else prod_fwd(pm_AAA, this->d_AA1, this->d_A1, this->d_AAA); + const auto &right = this->grace_model->AAA_right_A ? this->d_A : this->d_A1; // AAA = AA1 x A1 (legacy: x A) + if (shape_id) G1LF_SHAPE_FWD(AAA, this->d_AA1, right, this->d_AAA) + else prod_fwd(pm_AAA, this->d_AA1, right, this->d_AAA); if (this->fc_n_out == G1LF_NO) fc_fwd4(fcm_AA2, this->d_AA, this->d_A, this->d_AA2); else fc_fwd4(fcm_AA2, this->d_AA, this->d_A, this->d_AA2); if (shape_id) G1LF_SHAPE_FWD(AAAA, this->d_AA2, this->d_AA2, this->d_AAAA) @@ -817,7 +867,7 @@ void Grace1LFusedKK::stage_dense_fwd() } template -void Grace1LFusedKK::stage_reduce() { red_fwd<5>(); } +void Grace1LFusedKK::stage_reduce() { red_fwd(); } // the reverse reduce assigns A_adj .. AAAA_adj; A1_adj receives AAA (right) before AA (left, accumulated) template @@ -826,14 +876,16 @@ void Grace1LFusedKK::stage_zero_adj() { Kokkos::deep_copy(th template void Grace1LFusedKK::stage_dense_bwd() { - red_bwd<5, G1LF_NIN>(); + red_bwd(); if (shape_id) G1LF_SHAPE_REV(AAAA, this->d_AA2, this->d_AA2, this->d_AAAA_adj, this->d_AA2_adj, this->d_AA2_adj, (int) this->d_AA2_adj.extent(2)) else prod_rev(pm_AAAA, this->d_AA2, this->d_AA2, this->d_AAAA_adj, this->d_AA2_adj, this->d_AA2_adj, false); if (this->fc_n_out == G1LF_NO) fc_bwd(fcm_AA2, this->d_AA2_adj, this->d_AA_adj, this->d_A_adj); else fc_bwd(fcm_AA2, this->d_AA2_adj, this->d_AA_adj, this->d_A_adj); - if (shape_id) G1LF_SHAPE_REV(AAA, this->d_AA1, this->d_A1, this->d_AAA_adj, this->d_AA1_adj, this->d_A1_adj, ng_AAA) - else prod_rev(pm_AAA, this->d_AA1, this->d_A1, this->d_AAA_adj, this->d_AA1_adj, this->d_A1_adj, false); + const bool aaa_a = this->grace_model->AAA_right_A; + const auto &right = aaa_a ? this->d_A : this->d_A1, &right_adj = aaa_a ? this->d_A_adj : this->d_A1_adj; + if (shape_id) G1LF_SHAPE_REV(AAA, this->d_AA1, right, this->d_AAA_adj, this->d_AA1_adj, right_adj, ng_AAA) + else prod_rev(pm_AAA, this->d_AA1, right, this->d_AAA_adj, this->d_AA1_adj, right_adj, false); if (this->fc_n_out == G1LF_NO) fc_bwd(fcm_AA1, this->d_AA1_adj, this->d_AA_adj, this->d_A_adj); else fc_bwd(fcm_AA1, this->d_AA1_adj, this->d_AA_adj, this->d_A_adj); // AA reverse accumulates into A1_adj (generated with --acc AA); nothing to zero beyond the generated range diff --git a/libgrace/include/grace/kk/grace_1l_kk.h b/libgrace/include/grace/kk/grace_1l_kk.h index 730dd26..573774d 100644 --- a/libgrace/include/grace/kk/grace_1l_kk.h +++ b/libgrace/include/grace/kk/grace_1l_kk.h @@ -25,7 +25,7 @@ namespace kk { // Compile-time max dimensions for stack arrays in CUDA kernels static constexpr int GRACE1L_MAX_MLP_DIM = 64; // max hidden dim for MLP layers -static constexpr int GRACE1L_MAX_NRADMAX = 32; // max radial basis functions +static constexpr int GRACE1L_MAX_NRADMAX = 64; // max radial basis functions static constexpr int GRACE1L_MAX_MLP_LAYERS = 8; // max MLP depth static constexpr int GRACE1L_UQ_MAX_RP_DIM = 132; // bounds the FULL UQ feature dim D = rp_dim + n_density (z/f/delta stack arrays); 132 = 128 rp_dim + up to 3 density channels + 1 slack; a larger model fails the D > MAX guard at load @@ -71,7 +71,7 @@ class Grace1LKK { // ---- engine API (the LAMMPS glue and the Python module call these) ---- // model: loaded GRACE-1L model; type2mu[t] (t = 1..ntypes): model element of atom type t, -1 unmapped - void set_model(std::shared_ptr model, const std::vector &type2mu); + virtual void set_model(std::shared_ptr model, const std::vector &type2mu); virtual void setup(); // scalars, harmonics, device weights (LAMMPS: init_style) // another species set on a set-up engine: only the type tables (atom type -> model element, pair cutoffs and the // fused engines' per-type packed weights) are rebuilt; weights, spline tables and buffers stay @@ -421,6 +421,7 @@ class Grace1LKK { int energy_n_layers; int energy_max_dim; int energy_activation; // 0=silu, 1=tanh (device-side mirror) + bool energy_ln = false; // legacy models: RMS layer norm of the MLP input (scale folded into layer 0) // Per-element shifts t_nn_1d d_shifts; // [n_elements] diff --git a/libgrace/include/grace/kk/grace_1l_kk_impl.h b/libgrace/include/grace/kk/grace_1l_kk_impl.h index 2830625..cbb0ad3 100644 --- a/libgrace/include/grace/kk/grace_1l_kk_impl.h +++ b/libgrace/include/grace/kk/grace_1l_kk_impl.h @@ -430,6 +430,7 @@ void Grace1LKK::copy_weights_to_devi GRACE_KK_ERROR("GRACE-1L/KK: energy_mlp_activation={} not supported " "(only 0=silu, 1=tanh)", m.energy_mlp_activation); energy_activation = m.energy_mlp_activation; + energy_ln = m.energy_mlp_ln; int emax_in = 0, emax_out = 0; for (int i = 0; i < energy_n_layers; i++) { if (m.energy_mlp_layers[i].n_in > emax_in) emax_in = m.energy_mlp_layers[i].n_in; @@ -1898,6 +1899,10 @@ void Grace1LKK::operator() (TagCompu const int nin0 = d_energy_dims(0); for (int k = 0; k < nin0; k++) h_cur[k] = d_rho(ii, k + 1); + if (energy_ln) { // legacy models: RMS layer norm of the MLP input + const NNScalar q = inv_rms([&](int k) { return h_cur[k]; }, nin0); + for (int k = 0; k < nin0; k++) h_cur[k] *= q; + } // Hidden layers: 0 .. energy_n_layers-2 (silu or tanh activation) for (int layer = 0; layer < energy_n_layers - 1; layer++) { @@ -2780,6 +2785,9 @@ double Grace1LKK::cutoff(int mu_i, i template void Grace1LKK::setup() { + // the base stages evaluate the radial MLP with silu and AAA = AA1 x A1: legacy models run in the fused engine only + if (need_bond_buffers() && (grace_model->radial_activation || grace_model->AAA_right_A)) + GRACE_KK_ERROR("GRACE-1L/KK: legacy GRACE-1L model: the fused engine only"); // Copy scalar architecture params nelements = grace_model->n_elements; lmax = grace_model->lmax; @@ -2998,8 +3006,12 @@ Grace1LKK::run(const RunInput &in) int cs = chunk_size; int eact = energy_activation; const NNScalar os = output_scale; + const bool ln = energy_ln; Kokkos::parallel_for("ReverseMLPEnergy", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { + // legacy models: RMS layer norm z = x q of the MLP input x = rho[1:] + auto x = [&](int k) { return rho_in(ii, k + 1); }; + const NNScalar q = ln ? inv_rms(x, edims(0)) : NNScalar(1); // Stack arrays: pre_act for hidden layers, adj ping-pong NNScalar pre_act_buf[GRACE1L_MAX_MLP_LAYERS * GRACE1L_MAX_MLP_DIM]; NNScalar adj_a[GRACE1L_MAX_MLP_DIM], adj_b[GRACE1L_MAX_MLP_DIM]; @@ -3019,7 +3031,7 @@ Grace1LKK::run(const RunInput &in) NNScalar sum = 0.0; for (int k = 0; k < nin; k++) sum += eW(0, k, j) * rho_in(ii, k + 1); - pre_act[0 * emax + j] = sum * enorms(0); + pre_act[0 * emax + j] = (ln ? sum * q : sum) * enorms(0); } } // Hidden layers 1..en_layers-2: pre_act[layer] = norm * W @ act(pre_act[layer-1]) @@ -3088,8 +3100,9 @@ Grace1LKK::run(const RunInput &in) // Swap for next iteration NNScalar* tmp = adj_cur; adj_cur = adj_nxt; adj_nxt = tmp; } - // adj_cur now holds d_output/d_rho[k+1] for k=0..edims(0)-1 + // adj_cur now holds d_output/d_rho[k+1] for k=0..edims(0)-1 (legacy models: after the layer norm) const int nin0 = edims(0); + if (ln) ln_bwd([&](int k) -> NNScalar & { return adj_cur[k]; }, x, nin0, q); for (int k = 0; k < nin0; k++) rho_adj(ii, k + 1) = adj_cur[k]; }); diff --git a/libgrace/include/grace/kk/grace_2l_fused_kk.h b/libgrace/include/grace/kk/grace_2l_fused_kk.h index f81f5d0..f56b044 100644 --- a/libgrace/include/grace/kk/grace_2l_fused_kk.h +++ b/libgrace/include/grace/kk/grace_2l_fused_kk.h @@ -26,6 +26,23 @@ class Grace2LFusedKK : public Grace2LKK { // checks the model against the compiled kernels, builds the dense metadata and the radial tables void setup() override; + // a model without UQ narrower than a compiled width set runs zero-padded to it (GRACE2LModel::padded: the same + // results), with a warning + void set_model(std::shared_ptr model, const std::vector &type2mu) override; + // the compiled width set of a model (G2LF_WIDTHS); -1: outside the fused kernels + static int width_set(const GRACE2LModel &mm); + // the width set set_model pads a model to (the cheapest that fits); -1: none + static int pad_set(const GRACE2LModel &mm); + // the model zero-padded to its pad_set; the model itself without one + static std::shared_ptr padded_model(std::shared_ptr model); + // the warning for a model outside the compiled widths: zero-padded, or on the generic engine (use_generic); + // empty otherwise + static std::string slow_notice(const GRACE2LModel &mm); + // the model is outside the fused kernels, even padded: on the GPU only the generic engine Grace2LKK runs it + static bool needs_generic(const GRACE2LModel &mm) { return width_set(mm) < 0 && pad_set(mm) < 0; } + // the engine for a model: false this one, true the generic Grace2LKK (several + // times slower; warn with slow_notice): only with GRACE_GENERIC_KERNELS=1, an error (grace::Error) otherwise + static bool use_generic(const GRACE2LModel &mm); typename Base::RunOutput run(const typename Base::RunInput &in) override; void release_buffers() override { @@ -38,6 +55,7 @@ class Grace2LFusedKK : public Grace2LKK { int l1_nofit_nch = 0; // chunks of the run for which it did not fit (decided once per chunk count) int no_shape = 0; // 'noshape': force the generic CG-product kernels int spline_n = 1024; // cubic-Hermite intervals on x in [0, 1] ('spline' keyword) + int wset = 0; // setup: the compiled width set of the model (G2LF_WIDTHS) protected: std::vector st_A, st_A1, st_AA1, st_AA2; @@ -221,3 +239,5 @@ class Grace2LFusedKK : public Grace2LKK { #define GRACE_KK_EXTERN_2L(Device, Traits) \ extern template class grace::kk::Grace2LKK; \ extern template class grace::kk::Grace2LFusedKK +// Grace2LFusedKK::use_generic: a model outside the fused kernels runs on the generic engine Grace2LKK +#define GRACE_KK_2L_GENERIC_ENGINE 1 diff --git a/libgrace/include/grace/kk/grace_2l_fused_kk_impl.h b/libgrace/include/grace/kk/grace_2l_fused_kk_impl.h index 44e6374..2b8b1fb 100644 --- a/libgrace/include/grace/kk/grace_2l_fused_kk_impl.h +++ b/libgrace/include/grace/kk/grace_2l_fused_kk_impl.h @@ -6,6 +6,7 @@ #include "grace/kk/grace_2l_kk_impl.h" #include "grace/kk/fused_common.h" #include "grace/kk/fused_dense.h" +#include "grace/radial_x.h" #include "grace/kk/fused_hip_bonds.h" #include "grace/kk/g2l_shape_smaxl.h" #include "grace/kk/g2l_shape_smaxm.h" @@ -13,6 +14,7 @@ #include #include #include +#include #include #include @@ -30,11 +32,18 @@ constexpr int G2LF_NB2 = 16; // neighbours staged per batch, layer-2 forw constexpr int G2LF_MAX_YS = 4; // y-slots per thread in the dense L2 forward (nlm <= (TEAM / CB) * MAX_YS) constexpr int G2LF_MAX_B = 8; // per-thread B0 accumulators constexpr int G2LF_MAX_I = 16; // I_n_funcs (16 = l <= 3 message of the large models) -constexpr int G2LF_NO1 = 42; // FC widths compiled for (layer 1 / layer 2) +// the widths compiled for (GRACE2LModel::Widths): set 0 the current registry models, 1 the legacy GRACE-2L, 2 fits +// with 48 layer-2 radial channels; narrower models of either format run zero-padded to the cheapest set that fits +// (setup: wset). nq(): the float4 quads of the I1 / rho / I2 reducers. +struct G2LFWidths : GRACE2LModel::Widths { + constexpr int nq() const { return ((i1 > rho ? i1 : rho) + 3) / 4; } +}; +constexpr G2LFWidths G2LF_WIDTHS[] = {{{42, 64, 12, 17, 32}}, {{32, 32, 16, 13, 32}}, {{32, 48, 16, 13, 48}}}; +constexpr int G2LF_NSETS = std::size(G2LF_WIDTHS); +static_assert(G2LF_NSETS == 3, "G2LF_W dispatches three width sets"); // HIP: row stride of the layer-2 backward's M[y][channel][I] (I contiguous, 16-byte loads; 16 -> 20 floats so that // the 16-byte reads of consecutive channels do not conflict) constexpr int g2lf_nip(int nim) { return nim % 8 == 0 ? nim + 4 : nim; } -constexpr int G2LF_NO2 = 64; // Level-0 scratch of the vector-ALU layer-2 kernels for a channel block CB (the kernels' own formulas; gpu: their GPU // layouts, 16-byte rows). G2LF_SCRATCH: their budget (CUDA: below the 48 KB of a block); the GPU backward up to // G2LF_SCRATCH_BWD (HIP: the 64 KB of a CU; CUDA with Kokkos 5.2+: dynamic shared memory above 48 KB, 2L-large CB 16 @@ -105,28 +114,128 @@ using namespace g2lf_detail; /* ---------------------------------------------------------------------- */ +// the statement with W the model's compiled widths (G2LF_WIDTHS[wset]) +#define G2LF_W(...) do { switch (this->wset) { \ + case 0: { constexpr G2LFWidths W = G2LF_WIDTHS[0]; __VA_ARGS__; break; } \ + case 1: { constexpr G2LFWidths W = G2LF_WIDTHS[1]; __VA_ARGS__; break; } \ + case 2: { constexpr G2LFWidths W = G2LF_WIDTHS[2]; __VA_ARGS__; break; } \ + default: GRACE_KK_ERROR("GRACE-2L/KK: no compiled width set {}", this->wset); } } while (0) + +// the limits of the fused kernels besides the compiled widths, at no1 / nr2 layer-1 / layer-2 radial channels +inline bool g2lf_limits(const GRACE2LModel &mm, int no1, int nr2) +{ + const int nlm = (mm.lmax + 1) * (mm.lmax + 1); + return mm.lmax <= G2LF_MAX_LMAX && no1 * nlm <= G2LF_TEAM * G2LF_MAX_ACC && mm.I_n_funcs <= G2LF_MAX_I && + nlm <= (G2LF_TEAM / nr2) * G2LF_MAX_B; +} + +template +int Grace2LFusedKK::width_set(const GRACE2LModel &mm) +{ + if (mm.I_n_out != mm.L2_nradmax || mm.I2_n_out != mm.rho_n_out || !g2lf_limits(mm, mm.L1_nradmax, mm.L2_nradmax)) + return -1; + auto fc_width = [](std::initializer_list fcs, int no) { + for (const GRACE2LModel::FCWeights *w : fcs) + if (w->n_out != no || (int) (w->w_left.size() / ((size_t) w->n_out * w->w_shape_left)) != no) return false; + return true; + }; + const GRACE2LModel::Widths m = mm.widths(); + for (int s = 0; s < G2LF_NSETS; s++) { + const G2LFWidths &w = G2LF_WIDTHS[s]; + if (m.no1 == w.no1 && m.no2 == w.no2 && m.i1 == w.i1 && m.rho == w.rho && m.nr2 == w.nr2 && + fc_width({&mm.fc_A1, &mm.fc_AA1, &mm.fc_AA2}, w.no1) && fc_width({&mm.fc_B1, &mm.fc_BB1, &mm.fc_BB2}, w.no2)) + return s; + } + return -1; +} + +template +int Grace2LFusedKK::pad_set(const GRACE2LModel &mm) +{ + if (mm.has_uq || width_set(mm) >= 0) return -1; + const GRACE2LModel::Widths m = mm.widths(); + for (int s : {1, 2, 0}) { // cheapest first + const G2LFWidths &w = G2LF_WIDTHS[s]; + if (m.no1 <= w.no1 && m.no2 <= w.no2 && m.i1 <= w.i1 && m.rho <= w.rho && m.nr2 <= w.nr2 && + g2lf_limits(mm, w.no1, w.nr2)) + return s; + } + return -1; +} + +// GRACE_GENERIC_KERNELS=1: a model outside the fused kernels runs on the generic engine (read at every use_generic) +inline bool g2lf_generic_allowed() +{ + const char *v = std::getenv("GRACE_GENERIC_KERNELS"); + return v && std::atoi(v) != 0; +} + +// a model outside the compiled widths, for the warnings and the error +inline std::string g2lf_model(const GRACE2LModel &mm) +{ + return fmt("this GRACE-2L model (lmax {}, radial channels {} / {}, FC widths {} / {}, I1 {}, rho {})", mm.lmax, + mm.L1_nradmax, mm.L2_nradmax, mm.fc_A1.n_out, mm.fc_B1.n_out, mm.I1_n_out, mm.rho_n_out); +} + +template +std::string Grace2LFusedKK::slow_notice(const GRACE2LModel &mm) +{ + if (width_set(mm) >= 0) return ""; + if (const int s = pad_set(mm); s >= 0) { + const G2LFWidths &w = G2LF_WIDTHS[s]; + return fmt("{} has no GPU kernels of its own widths: it runs zero-padded to the compiled widths {} / {}, {}, {}, " + "{} (the same results, at the cost of the wider model). {}", g2lf_model(mm), w.no1, w.nr2, w.no2, w.i1, + w.rho, GRACE_KK_ISSUE); + } + if (!g2lf_generic_allowed()) return ""; + return g2lf_model(mm) + " has no fast GPU kernels: it runs on the generic GPU kernels (GRACE_GENERIC_KERNELS=1), " + "several times slower. " + GRACE_KK_ISSUE; +} + +template +bool Grace2LFusedKK::use_generic(const GRACE2LModel &mm) +{ + if (!needs_generic(mm)) return false; + if (!g2lf_generic_allowed()) + GRACE_KK_ERROR("GRACE-2L/KK: {} has no fast GPU kernels. Set GRACE_GENERIC_KERNELS=1 to run it on the generic GPU " + "kernels, several times slower. {}", g2lf_model(mm), GRACE_KK_ISSUE); + return true; +} + +template +std::shared_ptr Grace2LFusedKK::padded_model(std::shared_ptr model) +{ + const int s = model ? pad_set(*model) : -1; + if (s < 0) return model; + auto p = std::make_shared(model->padded(G2LF_WIDTHS[s])); + if (width_set(*p) != s) GRACE_KK_ERROR("GRACE-2L/KK: the padded model is outside width set {}", s); + return p; +} + +template +void Grace2LFusedKK::set_model(std::shared_ptr model, + const std::vector &type2mu) +{ + auto p = padded_model(model); + if (p != model) this->log.warn("GRACE-2L/KK: " + slow_notice(*model)); + Base::set_model(std::move(p), type2mu); +} + template void Grace2LFusedKK::setup() { if (spline_n < 16) GRACE_KK_ERROR("GRACE-2L/KK: spline N must be >= 16 (got {})", spline_n); + // the fused kernels are compiled for the widths of G2LF_WIDTHS (set_model pads narrower models); a model outside + // them runs on the generic engine Grace2LKK (use_generic) + const GRACE2LModel &mm = *this->grace_model; + wset = width_set(mm); + if (wset < 0) { // the caller did not route it to the generic engine (use_generic) + use_generic(mm); // without GRACE_GENERIC_KERNELS=1: the error naming it + GRACE_KK_ERROR("GRACE-2L/KK: {} has no fast GPU kernels, and this program does not run the generic GPU kernels " + "(update it)", g2lf_model(mm)); + } this->couple_supported = true; // two-group coupling in the bond kernels Base::setup(); // sets the model dimensions and device weights - if (this->lmax > G2LF_MAX_LMAX) - GRACE_KK_ERROR("GRACE-2L/KK: lmax {} > {} not supported", this->lmax, G2LF_MAX_LMAX); - if (this->L1_nradmax * (this->lmax + 1) * (this->lmax + 1) > G2LF_TEAM * G2LF_MAX_ACC) - GRACE_KK_ERROR("GRACE-2L/KK: A basis too large for the bond kernel"); - const int nlm = (this->lmax + 1) * (this->lmax + 1); - if (G2LF_TEAM % this->L2_nradmax != 0 || this->I_n_out != this->L2_nradmax || - this->L1_nradmax > G2LF_TEAM || this->I_n_funcs > G2LF_MAX_I || - nlm > (G2LF_TEAM / this->L2_nradmax) * G2LF_MAX_B) - GRACE_KK_ERROR("GRACE-2L/KK: model shape outside the compiled kernel limits"); - const GRACE2LModel &mm = *this->grace_model; - for (const GRACE2LModel::FCWeights *w : {&mm.fc_A1, &mm.fc_AA1, &mm.fc_AA2}) - if (w->n_out != G2LF_NO1 || (int) (w->w_left.size() / ((size_t) w->n_out * w->w_shape_left)) != G2LF_NO1) - GRACE_KK_ERROR("GRACE-2L/KK: layer-1 FC width must be {}", G2LF_NO1); - for (const GRACE2LModel::FCWeights *w : {&mm.fc_B1, &mm.fc_BB1, &mm.fc_BB2}) - if (w->n_out != G2LF_NO2 || (int) (w->w_left.size() / ((size_t) w->n_out * w->w_shape_left)) != G2LF_NO2) - GRACE_KK_ERROR("GRACE-2L/KK: layer-2 FC width must be {}", G2LF_NO2); // FC outputs that no CG product reads are not computed (forward) and have zero adjoint (reverse) auto live = [](const GRACE2LModel::CGProduct &p, bool left, bool right) { std::vector v; @@ -143,10 +252,6 @@ void Grace2LFusedKK::setup() this->log.info(fmt("[GRACE-2L/KK] live FC outputs A1 {}/{} AA1 {}/{} AA2 {}/{} B1 {}/{} BB1 {}/{} BB2 {}/{}\n", fcm_A1.nlive, fcm_A1.nf, fcm_AA1.nlive, fcm_AA1.nf, fcm_AA2.nlive, fcm_AA2.nf, fcm_B1.nlive, fcm_B1.nf, fcm_BB1.nlive, fcm_BB1.nf, fcm_BB2.nlive, fcm_BB2.nf)); - if (this->I1_n_out != 12 || this->B_n_out != G2LF_NO2 || this->L2_nradmax > 32 || this->L1_nradmax != G2LF_NO1) - GRACE_KK_ERROR("GRACE-2L/KK: ReduceN widths outside the compiled kernels (I1 12, B 64)"); - if (this->rho_n_out != 17) GRACE_KK_ERROR("GRACE-2L/KK: rho width must be 17"); - if (this->I2_n_out != 17) GRACE_KK_ERROR("GRACE-2L/KK: I2 width must be 17"); build_type_metas(); { std::vector sp = {{&mm.B_reduce_YI, 0, -1, this->B_n_out}, {&mm.B_reduce_B0, 1, -1, this->B_n_out}}; @@ -241,16 +346,18 @@ void Grace2LFusedKK::build_type_metas() const GRACE2LModel &mm = *this->grace_model; { const int nfo1 = this->I1_n_funcs; - std::vector sp = {{&mm.I1_reduce_A, 0, -1, 12}, {&mm.I1_reduce_AA, 1, -1, 12}, - {&mm.I1_reduce_AAA, 2, -1, 12}, {&mm.I1_reduce_AAAA, 3, -1, 12}, - {&mm.rho_reduce_A, 0, nfo1, 17}, {&mm.rho_reduce_AA, 1, nfo1, 17}, - {&mm.rho_reduce_AAA, 2, nfo1, 17}, {&mm.rho_reduce_AAAA, 3, nfo1, 17}}; + const int ni1 = mm.I1_n_out, nr = mm.rho_n_out; + std::vector sp = {{&mm.I1_reduce_A, 0, -1, ni1}, {&mm.I1_reduce_AA, 1, -1, ni1}, + {&mm.I1_reduce_AAA, 2, -1, ni1}, {&mm.I1_reduce_AAAA, 3, -1, ni1}, + {&mm.rho_reduce_A, 0, nfo1, nr}, {&mm.rho_reduce_AA, 1, nfo1, nr}, + {&mm.rho_reduce_AAA, 2, nfo1, nr}, {&mm.rho_reduce_AAAA, 3, nfo1, nr}}; build_red_meta(rm_I1, sp, nfo1 + 1, nfo1, true, {(int) this->n_funcs_A1, (int) this->n_funcs_AA, (int) this->n_funcs_AAA, (int) this->n_funcs_AAAA}); } { - std::vector sp = {{&mm.I2_reduce_B, 0, 0, 17}, {&mm.I2_reduce_BB, 1, 0, 17}, - {&mm.I2_reduce_BBB, 2, 0, 17}, {&mm.I2_reduce_BBBB, 3, 0, 17}}; + const int nr = mm.I2_n_out; + std::vector sp = {{&mm.I2_reduce_B, 0, 0, nr}, {&mm.I2_reduce_BB, 1, 0, nr}, + {&mm.I2_reduce_BBB, 2, 0, nr}, {&mm.I2_reduce_BBBB, 3, 0, nr}}; build_red_meta(rm_I2, sp, 1, 0, true, {(int) this->B_n_funcs, (int) this->n_funcs_BB, (int) this->n_funcs_BBB, (int) this->n_funcs_BBBB}); } @@ -268,7 +375,7 @@ void Grace2LFusedKK::build_spline_tables() std::vector g, dg, out, dout; for (int k = 0; k <= N; k++) { radial_basis_x(k * h, p, nrb, g, dg); - mlp_x(layers, g, dg, out, dout); + mlp_x(layers, g, dg, out, dout, m.radial_activation); for (int o = 0; o < nout; o++) { fv[(size_t) k * nout + o] = out[o]; dv[(size_t) k * nout + o] = h * dout[o]; } } dc = V3(std::string("g2lf:") + name, N, nout, 4); @@ -650,8 +757,8 @@ int Grace2LFusedKK::l2_channel_block(bool bwd, int *rounds) return ys <= G2LF_MAX_YS ? g2lf_fwd_t_shmem(gpu, nim, NI, nl, CB) : SIZE_MAX; }; auto widest = [&](int r) { // 0: none fits - for (int CB = NR2; CB >= 8; CB /= 2) - if (NR2 % CB == 0 && G2LF_TEAM % CB == 0 && shmem(r, CB) <= budget) return CB; + for (int CB = 32; CB >= 8; CB /= 2) // powers of two (they divide the team) that divide NR2 + if (CB <= NR2 && NR2 % CB == 0 && shmem(r, CB) <= budget) return CB; return 0; }; // backward, GPU teams of 256: the geometry of 3 rounds per step; one round where no channel block fits with 3 @@ -2014,23 +2121,23 @@ bool Grace2LFusedKK::fc_forward_own(typename Base::FCLayer w using L = typename Base::FCLayer; const FCMeta &m = which == L::A1 ? fcm_A1 : which == L::AA1 ? fcm_AA1 : which == L::AA2 ? fcm_AA2 : which == L::B1 ? fcm_B1 : which == L::BB1 ? fcm_BB1 : fcm_BB2; - if (m.no == this->L1_fc_n_out) fc_fwd(m, in_left, in_right, out); - else fc_fwd4(m, in_left, in_right, out); + if (which == L::A1 || which == L::AA1 || which == L::AA2) G2LF_W(fc_fwd(m, in_left, in_right, out)); + else G2LF_W(fc_fwd4(m, in_left, in_right, out)); return true; } template -void Grace2LFusedKK::stage_rev_fc_AA2() { fc_bwd(fcm_AA2, this->d_AA2_adj, this->d_AA_adj, this->d_A_adj); } +void Grace2LFusedKK::stage_rev_fc_AA2() { G2LF_W(fc_bwd(fcm_AA2, this->d_AA2_adj, this->d_AA_adj, this->d_A_adj)); } template -void Grace2LFusedKK::stage_rev_fc_AA1() { fc_bwd(fcm_AA1, this->d_AA1_adj, this->d_AA_adj, this->d_A_adj); } +void Grace2LFusedKK::stage_rev_fc_AA1() { G2LF_W(fc_bwd(fcm_AA1, this->d_AA1_adj, this->d_AA_adj, this->d_A_adj)); } template -void Grace2LFusedKK::stage_rev_fc_A1() { fc_bwd(fcm_A1, this->d_A1_adj, this->d_A_adj, this->d_A_adj); } +void Grace2LFusedKK::stage_rev_fc_A1() { G2LF_W(fc_bwd(fcm_A1, this->d_A1_adj, this->d_A_adj, this->d_A_adj)); } template -void Grace2LFusedKK::stage_rev_fc_BB2() { fc_bwd(fcm_BB2, this->d_BB2_adj, this->d_BB_adj, this->d_B_adj); } +void Grace2LFusedKK::stage_rev_fc_BB2() { G2LF_W(fc_bwd(fcm_BB2, this->d_BB2_adj, this->d_BB_adj, this->d_B_adj)); } template -void Grace2LFusedKK::stage_rev_fc_BB1() { fc_bwd(fcm_BB1, this->d_BB1_adj, this->d_BB_adj, this->d_B_adj); } +void Grace2LFusedKK::stage_rev_fc_BB1() { G2LF_W(fc_bwd(fcm_BB1, this->d_BB1_adj, this->d_BB_adj, this->d_B_adj)); } template -void Grace2LFusedKK::stage_rev_fc_B1() { fc_bwd(fcm_B1, this->d_B1_adj, this->d_B_adj, this->d_B_adj); } +void Grace2LFusedKK::stage_rev_fc_B1() { G2LF_W(fc_bwd(fcm_B1, this->d_B1_adj, this->d_B_adj, this->d_B_adj)); } /* ---------------------------------------------------------------------- CG products in gather form (one write per element, no per-term read-modify-write): @@ -2414,21 +2521,21 @@ void Grace2LFusedKK::red_bwd(const RedMeta &m, const t_nn3 & } template void Grace2LFusedKK::stage_reduce_B() -{ red_fwd(rm_B, this->d_YI, this->d_B0, this->d_YI, this->d_YI, this->d_B, this->d_rho); } +{ G2LF_W(red_fwd<(W.no2 + 3) / 4, 0, W.no2, 0>(rm_B, this->d_YI, this->d_B0, this->d_YI, this->d_YI, this->d_B, this->d_rho)); } template void Grace2LFusedKK::stage_rev_reduce_B() -{ red_bwd(rm_B, this->d_B_adj, this->d_rho_adj, this->d_YI_adj, this->d_B0_adj, this->d_YI_adj, this->d_YI_adj); } +{ G2LF_W(red_bwd<(W.no2 + 3) / 4, W.nr2, 0, W.no2, 0, true>(rm_B, this->d_B_adj, this->d_rho_adj, this->d_YI_adj, this->d_B0_adj, this->d_YI_adj, this->d_YI_adj)); } // the I1 / rho reduction (rho_only: rho alone, same blocks and order) template void Grace2LFusedKK::reduce_I1(bool rho_only) { auto A = this->d_A; auto AA = this->d_AA; auto AAA = this->d_AAA; auto AAAA = this->d_AAAA; if (red_blk_on(rm_I1)) - red_fwd<5, 1, 12, 17, true>(rm_I1, A, AA, AAA, AAAA, this->d_I1, this->d_rho, rho_only); - else red_fwd<5, 1, 12, 17>(rm_I1, A, AA, AAA, AAAA, this->d_I1, this->d_rho, rho_only); + G2LF_W(red_fwd(rm_I1, A, AA, AAA, AAAA, this->d_I1, this->d_rho, rho_only)); + else G2LF_W(red_fwd(rm_I1, A, AA, AAA, AAAA, this->d_I1, this->d_rho, rho_only)); } template void Grace2LFusedKK::stage_reduce_I1() { reduce_I1(false); } template void Grace2LFusedKK::stage_rev_reduce_I1() -{ red_bwd<5, G2LF_NO1, 1, 12, 17, true>(rm_I1, this->d_I1_adj, this->d_rho_adj, this->d_A_adj, this->d_AA_adj, this->d_AAA_adj, this->d_AAAA_adj); } +{ G2LF_W(red_bwd(rm_I1, this->d_I1_adj, this->d_rho_adj, this->d_A_adj, this->d_AA_adj, this->d_AAA_adj, this->d_AAAA_adj)); } template void Grace2LFusedKK::stage_reduce_rho() {} // done in stage_reduce_I1 // rho of stage_reduce_I1 alone (the base kernel sums in another fp32 order: forces about 1e-6 relative off the energy // gradient in multi-chunk runs without the layer-1 store) @@ -2442,11 +2549,11 @@ template void Grace2LFusedKK // I2 <- B, BB, BBB, BBBB (scalar outputs, element-dependent): every connection to the single output fo = 0 template void Grace2LFusedKK::stage_reduce_I2() { - if (red_blk_on(rm_I2)) red_fwd<5, 1, 0, 17, true>(rm_I2, this->d_B, this->d_BB, this->d_BBB, this->d_BBBB, this->d_B, this->d_I2); - else red_fwd<5, 1, 0, 17>(rm_I2, this->d_B, this->d_BB, this->d_BBB, this->d_BBBB, this->d_B, this->d_I2); + if (red_blk_on(rm_I2)) G2LF_W(red_fwd(rm_I2, this->d_B, this->d_BB, this->d_BBB, this->d_BBBB, this->d_B, this->d_I2)); + else G2LF_W(red_fwd(rm_I2, this->d_B, this->d_BB, this->d_BBB, this->d_BBBB, this->d_B, this->d_I2)); } template void Grace2LFusedKK::stage_rev_reduce_I2() -{ red_bwd<5, G2LF_NO2, 1, 0, 17, true>(rm_I2, this->d_B_adj, this->d_I2_adj, this->d_B_adj, this->d_BB_adj, this->d_BBB_adj, this->d_BBBB_adj); } +{ G2LF_W(red_bwd(rm_I2, this->d_B_adj, this->d_I2_adj, this->d_B_adj, this->d_BB_adj, this->d_BBB_adj, this->d_BBBB_adj)); } } // namespace kk } // namespace grace diff --git a/libgrace/include/grace/kk/grace_2l_kk.h b/libgrace/include/grace/kk/grace_2l_kk.h index d1b623e..d25300e 100644 --- a/libgrace/include/grace/kk/grace_2l_kk.h +++ b/libgrace/include/grace/kk/grace_2l_kk.h @@ -25,7 +25,7 @@ namespace kk { static constexpr int GRACE2L_MAX_MLP_DIM = 210; // max hidden/output dim for MLP layers static constexpr int GRACE2L_MAX_MLP_HIDDEN = 128; // max input/hidden dim for MLP radial buffers (NOT output) -static constexpr int GRACE2L_MAX_NRADMAX = 42; // max radial basis functions +static constexpr int GRACE2L_MAX_NRADMAX = 64; // max radial basis functions (the fused engine: its width sets) static constexpr int GRACE2L_MAX_MLP_LAYERS = 8; // max MLP depth static constexpr int GRACE2L_MAX_YI_OUT_FUNCS = 200; // max YI output functions static constexpr int GRACE2L_MAX_I1_FUNCS = 16; // max I1 equivariant channels (I1_n_funcs) @@ -80,7 +80,7 @@ class Grace2LKK { // ---- engine API (the LAMMPS glue and the Python module call these) ---- // model: loaded GRACE-2L model; type2mu[t] (t = 1..ntypes): model element of atom type t, -1 unmapped - void set_model(std::shared_ptr model, const std::vector &type2mu); + virtual void set_model(std::shared_ptr model, const std::vector &type2mu); virtual void setup(); // scalars, harmonics, device weights, cuBLAS (LAMMPS: init_style) // another species set on a set-up engine: only the type tables (atom type -> model element, pair cutoffs and the // fused engines' per-type packed weights) are rebuilt; weights, spline tables and buffers stay @@ -722,6 +722,10 @@ class Grace2LKK { int energy_n_layers; int energy_max_dim; int energy_activation; // 0=silu, 1=tanh + int rms_n1 = 0, rms_n0 = 0; // the RMS means of rho[1:] / I2 divide by these (padded model: the unpadded widths) + bool energy_rms = true; // the RMSNorms I_nl_LN, I_0_LN (off: legacy models, unit scales) + int radial_activation = 0; // radial MLPs R, R1: 0=silu, 1=tanh (legacy models) + bool energy_ln = false; // legacy models: RMS layer norm of the MLP input (scale folded into layer 0) // Per-element shifts t_nn_1d d_shifts; @@ -878,6 +882,20 @@ class Grace2LKK { return Kokkos::tanh(x); } + // a hidden unit of the radial MLPs in forward mode: h = act(s), dh = act'(s) ds (act 0 silu, 1 tanh: legacy models) + KOKKOS_INLINE_FUNCTION + static void radial_act(int act, NNScalar s, NNScalar ds, NNScalar &h, NNScalar &dh) { + const NNScalar one(1); + if (act == 1) { + h = Kokkos::tanh(s); + dh = (one - h * h) * ds; + return; + } + const NNScalar sig = one / (one + Kokkos::exp(-s)); + h = s * sig; + dh = sig * (one + s * (one - sig)) * ds; + } + // ---------- UQ / extrapolation grade (forward only) ---------- // FP64 device storage regardless of NNScalar (numerical safety). diff --git a/libgrace/include/grace/kk/grace_2l_kk_impl.h b/libgrace/include/grace/kk/grace_2l_kk_impl.h index 53b3d59..86024fb 100644 --- a/libgrace/include/grace/kk/grace_2l_kk_impl.h +++ b/libgrace/include/grace/kk/grace_2l_kk_impl.h @@ -718,6 +718,11 @@ void Grace2LKK::copy_weights_to_devi GRACE_KK_ERROR("GRACE-2L/KK: energy_mlp_activation={} not supported " "(only 0=silu, 1=tanh)", m.energy_mlp_activation); energy_activation = m.energy_mlp_activation; + energy_rms = m.energy_rms_norm; + rms_n1 = m.I_nl_LN_n_out - 1; + rms_n0 = m.I_0_LN_n_out; + radial_activation = m.radial_activation; + energy_ln = m.energy_mlp_ln; copy_mlp(m.energy_mlp_layers, energy_n_layers, d_energy_W, d_energy_norms, d_energy_dims, energy_max_dim, "g2l:energy"); if (energy_n_layers > GRACE2L_MAX_MLP_LAYERS) @@ -1312,7 +1317,8 @@ Grace2LKK::run(const RunInput &in) auto lnout = d_I_nl_LN; auto scale = d_I_nl_LN_scale; auto I_nl_glob = d_I_nl_LN_global; - int n_out_ln = rho_n_out; + int n_out_ln = rho_n_out, nrms = rms_n1; + const bool rms = energy_rms; Kokkos::parallel_for("RMSNorm_I_nl_LN", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { @@ -1324,7 +1330,7 @@ Grace2LKK::run(const RunInput &in) NNScalar sq_sum = 0.0; for (int k = 1; k < n_out_ln; k++) sq_sum += rho_in(ii, k) * rho_in(ii, k); - NNScalar rms_inv = 1.0 / Kokkos::sqrt(sq_sum / (n_out_ln - 1) + 1e-8); + NNScalar rms_inv = rms ? 1.0 / Kokkos::sqrt(sq_sum / nrms + 1e-8) : 1.0; for (int k = 1; k < n_out_ln; k++) { NNScalar val = rho_in(ii, k) * rms_inv * scale(k - 1); lnout(ii, k) = val; @@ -1493,14 +1499,15 @@ Grace2LKK::run(const RunInput &in) auto I2_in = d_I2; auto lnout = d_I_0_LN; auto scale = d_I_0_LN_scale; - int n_out_ln = I2_n_out, cs = chunk_size; + int n_out_ln = I2_n_out, nrms = rms_n0, cs = chunk_size; + const bool rms = energy_rms; Kokkos::parallel_for("RMSNorm_I_0_LN", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { NNScalar sq_sum = 0.0; for (int k = 0; k < n_out_ln; k++) sq_sum += I2_in(ii, k) * I2_in(ii, k); - NNScalar rms_inv = 1.0 / Kokkos::sqrt(sq_sum / n_out_ln + 1e-8); + NNScalar rms_inv = rms ? 1.0 / Kokkos::sqrt(sq_sum / nrms + 1e-8) : 1.0; for (int k = 0; k < n_out_ln; k++) lnout(ii, k) = I2_in(ii, k) * rms_inv * scale(k); }); @@ -1577,6 +1584,7 @@ Grace2LKK::run(const RunInput &in) auto il = d_ilist; int co = chunk_offset; const NNScalar os = output_scale; auto wk = d_emlp_work; // per atom: pre-activations [en_layers - 1][emax], then two adjoint rows [emax] + const bool ln = energy_ln; Kokkos::parallel_for("ReverseMLPEnergy", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { @@ -1584,18 +1592,20 @@ Grace2LKK::run(const RunInput &in) // as in the interleaved local memory it replaces) int adj_cur = (en_layers - 1) * emax, adj_nxt = adj_cur + emax; // column offsets of the two rows - // Input: I_nl_LN[:,1:] + I_0_LN[:,1:] + // Input x = I_nl_LN[:,1:] + I_0_LN[:,1:] (legacy models: its RMS layer norm z = x q) const int nin0 = edims(0); + const int i = il[ii + co]; + auto x = [&](int k) { return I_nl_in(i, k + 1) + I_0_in(ii, k + 1); }; + const NNScalar q = ln ? inv_rms(x, nin0) : NNScalar(1); // Forward pass: recompute hidden pre-activations { const int nout = edims(1); for (int j = 0; j < nout; j++) { NNScalar sum = 0.0; - const int i = il[ii + co]; for (int k = 0; k < nin0; k++) sum += eW(0, k, j) * (I_nl_in(i, k + 1) + I_0_in(ii, k + 1)); - wk(ii, 0 * emax + j) = sum * enorms(0); + wk(ii, 0 * emax + j) = (ln ? sum * q : sum) * enorms(0); } } for (int layer = 1; layer < en_layers - 1; layer++) { @@ -1647,7 +1657,7 @@ Grace2LKK::run(const RunInput &in) // adj_cur now holds d_output/d_input[k] for k=0..nin0-1 // Both I_nl_LN and I_0_LN get the same gradient (sum rule). // Linear skip seeds scaled by output_scale (chain rule for E*output_scale). - const int i = il[ii + co]; + if (ln) ln_bwd([&](int k) -> NNScalar & { return wk(ii, adj_cur + k); }, x, nin0, q); nl_adj(ii, 0) = os; // linear skip term, scaled i0_adj(ii, 0) = os; nl_adj_glob(i, 0) = os; @@ -1664,7 +1674,8 @@ Grace2LKK::run(const RunInput &in) auto I2_in = d_I2; auto ln_adj = d_I_0_LN_adj; auto ln_out = d_I_0_LN; auto I2_adj_out = d_I2_adj; auto scale = d_I_0_LN_scale; - int n_out_ln = I2_n_out, cs = chunk_size; + int n_out_ln = I2_n_out, nrms = rms_n0, cs = chunk_size; + const bool rms = energy_rms; Kokkos::parallel_for("ReverseRMSNorm_I_0_LN", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { @@ -1672,16 +1683,17 @@ Grace2LKK::run(const RunInput &in) NNScalar sq_sum = 0.0; for (int k = 0; k < n_out_ln; k++) sq_sum += I2_in(ii, k) * I2_in(ii, k); - NNScalar rms_inv = 1.0 / Kokkos::sqrt(sq_sum / n_out_ln + 1e-8); + NNScalar rms_inv = rms ? 1.0 / Kokkos::sqrt(sq_sum / nrms + 1e-8) : 1.0; - // dot = Σ_k d_ln[k] * ln_out[k] + // dot = Σ_k d_ln[k] * ln_out[k] (without the RMSNorm: 0, the identity) NNScalar dot = 0.0; - for (int k = 0; k < n_out_ln; k++) - dot += ln_adj(ii, k) * ln_out(ii, k); + if (rms) + for (int k = 0; k < n_out_ln; k++) + dot += ln_adj(ii, k) * ln_out(ii, k); // dx[k] = rms_inv * (d_ln[k] * scale[k] - x[k] * dot * rms_inv / N) for (int k = 0; k < n_out_ln; k++) - I2_adj_out(ii, k) = rms_inv * (ln_adj(ii, k) * scale(k) - I2_in(ii, k) * dot * rms_inv / n_out_ln); + I2_adj_out(ii, k) = rms_inv * (ln_adj(ii, k) * scale(k) - I2_in(ii, k) * dot * rms_inv / nrms); }); } @@ -1840,8 +1852,9 @@ Grace2LKK::run(const RunInput &in) auto I_nl_out = d_I_nl_LN_global; auto rho_adj_out = d_rho_adj; auto scale = d_I_nl_LN_scale; - int n_out_ln = rho_n_out, cs = chunk_size; + int n_out_ln = rho_n_out, nrms = rms_n1, cs = chunk_size; auto il = d_ilist; int co = chunk_offset; + const bool rms = energy_rms; Kokkos::parallel_for("ReverseRMSNorm_I_nl_LN", Kokkos::RangePolicy(0, cs), KOKKOS_LAMBDA(const int ii) { @@ -1853,15 +1866,16 @@ Grace2LKK::run(const RunInput &in) NNScalar sq_sum = 0.0; for (int k = 1; k < n_out_ln; k++) sq_sum += rho_in(ii, k) * rho_in(ii, k); - NNScalar rms_inv = 1.0 / Kokkos::sqrt(sq_sum / (n_out_ln - 1) + 1e-8); + NNScalar rms_inv = rms ? 1.0 / Kokkos::sqrt(sq_sum / nrms + 1e-8) : 1.0; - // dot = Σ_{k=1}^{n-1} d_ln[k] * ln_out[k] + // dot = Σ_{k=1}^{n-1} d_ln[k] * ln_out[k] (without the RMSNorm: 0, the identity) NNScalar dot = 0.0; - for (int k = 1; k < n_out_ln; k++) - dot += nl_adj(i, k) * I_nl_out(i, k); + if (rms) + for (int k = 1; k < n_out_ln; k++) + dot += nl_adj(i, k) * I_nl_out(i, k); for (int k = 1; k < n_out_ln; k++) - rho_adj_out(ii, k) = rms_inv * (nl_adj(i, k) * scale(k - 1) - rho_in(ii, k) * dot * rms_inv / (n_out_ln - 1)); + rho_adj_out(ii, k) = rms_inv * (nl_adj(i, k) * scale(k - 1) - rho_in(ii, k) * dot * rms_inv / nrms); }); } @@ -2051,10 +2065,7 @@ void Grace2LKK::operator() (TagCompu NNScalar w = d_mlp_rad_R_W(layer, k, j); s += w * h_cur[k]; ds += w * dh_cur[k]; } - s *= norm; ds *= norm; - NNScalar sig = 1.0 / (1.0 + Kokkos::exp(-s)); - h_nxt[j] = s * sig; - dh_nxt[j] = sig * (1.0 + s * (1.0 - sig)) * ds; + radial_act(radial_activation, s * norm, ds * norm, h_nxt[j], dh_nxt[j]); } NNScalar* tmp; tmp = h_cur; h_cur = h_nxt; h_nxt = tmp; @@ -2139,10 +2150,7 @@ void Grace2LKK::operator() (TagCompu NNScalar w = d_mlp_rad_R1_W(layer, k, j); s += w * h_cur[k]; ds += w * dh_cur[k]; } - s *= norm; ds *= norm; - NNScalar sig = 1.0 / (1.0 + Kokkos::exp(-s)); - h_nxt[j] = s * sig; - dh_nxt[j] = sig * (1.0 + s * (1.0 - sig)) * ds; + radial_act(radial_activation, s * norm, ds * norm, h_nxt[j], dh_nxt[j]); } NNScalar* tmp; tmp = h_cur; h_cur = h_nxt; h_nxt = tmp; @@ -2218,8 +2226,8 @@ void Grace2LKK::operator() (TagMLPAs // TagMLPActDeriv: fused value+deriv silu applied in place between GEMMs. Reads // the raw value GEMM output (d_mlp_h{0,1}) and raw deriv GEMM output // (d_mlp_dh{0,1}); forms s = raw*norm (2L MLP has NO bias) and ds = raw_deriv*norm, -// then writes silu(s) and silu'(s)*ds. silu'(s) = sig*(1 + s*(1-sig)) -- identical -// to the hand kernel TagComputeMLPRadial_R{,1}. +// then writes silu(s) and silu'(s)*ds (legacy models: tanh), as the hand kernel TagComputeMLPRadial_R{,1} +// (radial_act). template KOKKOS_INLINE_FUNCTION void Grace2LKK::operator() (TagMLPActDeriv, @@ -2229,19 +2237,12 @@ void Grace2LKK::operator() (TagMLPAc const int b = (int) (idx / nout); const int j = (int) (idx - (int64_t) b * nout); const NNScalar norm = mlp_act_norm; - if (mlp_act_layer == 0) { - const NNScalar s = d_mlp_h0(b, j) * norm; - const NNScalar ds = d_mlp_dh0(b, j) * norm; - const NNScalar sig = NNScalar(1.0) / (NNScalar(1.0) + Kokkos::exp(-s)); - d_mlp_h0(b, j) = s * sig; - d_mlp_dh0(b, j) = sig * (NNScalar(1.0) + s * (NNScalar(1.0) - sig)) * ds; - } else { - const NNScalar s = d_mlp_h1(b, j) * norm; - const NNScalar ds = d_mlp_dh1(b, j) * norm; - const NNScalar sig = NNScalar(1.0) / (NNScalar(1.0) + Kokkos::exp(-s)); - d_mlp_h1(b, j) = s * sig; - d_mlp_dh1(b, j) = sig * (NNScalar(1.0) + s * (NNScalar(1.0) - sig)) * ds; - } + const auto &h = mlp_act_layer == 0 ? d_mlp_h0 : d_mlp_h1; + const auto &dh = mlp_act_layer == 0 ? d_mlp_dh0 : d_mlp_dh1; + NNScalar hv, dhv; + radial_act(radial_activation, h(b, j) * norm, dh(b, j) * norm, hv, dhv); + h(b, j) = hv; + dh(b, j) = dhv; } // TagMLPScatterR: write the LayoutRight output-GEMM result d_mlp_R[b, j] (already @@ -3150,6 +3151,10 @@ void Grace2LKK::operator() (TagCompu const int nin0 = d_energy_dims(0); for (int k = 0; k < nin0; k++) h_cur[k] = d_I_nl_LN_global(i, k + 1) + d_I_0_LN(ii, k + 1); + if (energy_ln) { // legacy models: RMS layer norm of the MLP input + const NNScalar q = inv_rms([&](int k) { return h_cur[k]; }, nin0); + for (int k = 0; k < nin0; k++) h_cur[k] *= q; + } // Hidden layers with activation for (int layer = 0; layer < energy_n_layers - 1; layer++) { diff --git a/libgrace/include/grace/kk/support.h b/libgrace/include/grace/kk/support.h index fb7e11b..292fcf3 100644 --- a/libgrace/include/grace/kk/support.h +++ b/libgrace/include/grace/kk/support.h @@ -266,6 +266,22 @@ template KOKKOS_INLINE_FUNCTION BondSlot bond_slot(const Member &t return {team.team_rank() + ts * (lr % teams_per_slot), lr / teams_per_slot}; } +// Legacy GRACE-1L/2L: the RMS layer norm z = x q of the energy-MLP input (its scale folded into the first layer), q = +// 1 / sqrt(mean x^2 + 1e-16) as tensorpotential's scalar_rms_ln. inv_rms: q of x(0..n); ln_bwd: g(k) = dE/dz_k -> dE/dx_k. +template KOKKOS_INLINE_FUNCTION T inv_rms(const X &x, int n) +{ + T ms = 0; + for (int k = 0; k < n; k++) ms += x(k) * x(k); + return T(1) / Kokkos::sqrt(ms / n + T(1e-16)); +} + +template KOKKOS_INLINE_FUNCTION void ln_bwd(const G &g, const X &x, int n, T q) +{ + T t = 0; + for (int k = 0; k < n; k++) t += g(k) * x(k); + for (int k = 0; k < n; k++) g(k) = q * (g(k) - x(k) * q * q * t / n); +} + // A RangePolicy with 64-bit indices, for ranges that can pass 2^31 (bonds of a whole system). Kokkos >= 5.2 indexes a // RangePolicy without IndexType with int on CUDA (before: unsigned int) and aborts the process on a larger bound. template using Range64 = Kokkos::RangePolicy>; diff --git a/libgrace/include/grace/model_1l.h b/libgrace/include/grace/model_1l.h index 84714ab..3d6c762 100644 --- a/libgrace/include/grace/model_1l.h +++ b/libgrace/include/grace/model_1l.h @@ -87,6 +87,9 @@ struct GRACE1LModel { // Energy MLP std::vector energy_mlp_layers; int energy_mlp_activation = 0; // 0=silu, 1=tanh + int radial_activation = 0; // radial MLP: 0=silu, 1=tanh (legacy models) + bool AAA_right_A = false; // legacy models: AAA = AA1 x A (current: AA1 x A1) + bool energy_mlp_ln = false; // legacy models: RMS layer norm of the MLP input, x s / rms(x) (s folded into layer 0) // Per-element shifts std::vector shift_values; @@ -123,6 +126,10 @@ struct GRACE1LModel { // load() of an archive already read (key -> array), e.g. a modified copy in a test void load_arrays(const std::map &npz); + // a copy zero-padded to wider channel widths, with the same results (models without UQ only): radial channels + // (A) nr, FC outputs no + GRACE1LModel padded(int nr, int no) const; + // throws grace::Error unless the model carries a usable schema-6 UQ head void validate_uq() const; diff --git a/libgrace/include/grace/model_2l.h b/libgrace/include/grace/model_2l.h index 1d800fe..f294a8d 100644 --- a/libgrace/include/grace/model_2l.h +++ b/libgrace/include/grace/model_2l.h @@ -101,7 +101,7 @@ struct GRACE2LModel { // RMSNorm: I_nl_LN (only_nonlin: scale has n_out-1 elements, rho[0] passed through) std::vector I_nl_LN_scale; // [rho_n_out - 1] = [16] int I_nl_LN_type = 1; // 1 = only_nonlin - int I_nl_LN_n_out = 17; + int I_nl_LN_n_out = 17; // the RMS mean runs over rho[1:I_nl_LN_n_out] (padded(): the unpadded rho width) // ---- Layer 2 ---- int L2_nradmax = 32; @@ -141,11 +141,15 @@ struct GRACE2LModel { // RMSNorm: I_0_LN (full: scale has n_out elements, all scaled) std::vector I_0_LN_scale; // [I2_n_out] = [17] int I_0_LN_type = 0; // 0 = full - int I_0_LN_n_out = 17; + int I_0_LN_n_out = 17; // the RMS mean runs over I2[:I_0_LN_n_out] (padded(): the unpadded I2 width) // Energy MLP (input=16, tanh activation) std::vector energy_mlp_layers; int energy_mlp_activation = 1; // 0=silu, 1=tanh + int radial_activation = 0; // radial MLPs R, R1: 0=silu, 1=tanh (legacy models) + bool energy_mlp_ln = false; // legacy models: RMS layer norm of the MLP input, x s / rms(x) (s folded into layer 0) + // false (legacy models): rho and I2 enter the energy MLP as they are (unit scales, the engines skip the 1/rms) + bool energy_rms_norm = true; // Per-element shifts std::vector shift_values; @@ -182,6 +186,15 @@ struct GRACE2LModel { // load() of an archive already read (key -> array), e.g. a modified copy in a test void load_arrays(const std::map &npz); + // channel widths: layer-1 radial and FC channels no1, layer-2 FC and B channels no2, I1, rho = I2, layer-2 radial + // channels = I nr2 + struct Widths { + int no1, no2, i1, rho, nr2; + }; + Widths widths() const { return {L1_nradmax, B_n_out, I1_n_out, rho_n_out, L2_nradmax}; } + // a copy zero-padded to wider channel widths, with the same results (models without UQ only) + GRACE2LModel padded(const Widths &w) const; + // throws grace::Error unless the model carries a usable schema-6 UQ head void validate_uq() const; diff --git a/libgrace/include/grace/radial_x.h b/libgrace/include/grace/radial_x.h new file mode 100644 index 0000000..d5984ab --- /dev/null +++ b/libgrace/include/grace/radial_x.h @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: LicenseRef-ASL (Academic Software Licence, see LICENSE.md) +// Host fp64 radial functions in forward mode (value and d/dx), for the spline tables of the CPU and Kokkos engines +#pragma once + +#include +#include + +namespace grace { + +// Chebyshev T_1..T_nrb of (2x-1) times fcut_p(x), and d/dx (x = r / rcut_ij) +inline void radial_basis_x(double x, int p, int nrb, std::vector &g, std::vector &dg) +{ + const double xp = std::pow(x, p), xp1 = xp * x, xp2 = xp1 * x; + const double fc = 1.0 - 0.5 * (p + 1) * (p + 2) * xp + p * (p + 2) * xp1 - 0.5 * p * (p + 1) * xp2; + const double dfc = 0.5 * p * (p + 1) * (p + 2) * (-std::pow(x, p - 1) + 2.0 * xp - xp1); + const double t = 2.0 * x - 1.0; + double Tp = 1.0, Tc = t, dTp = 0.0, dTc = 2.0; + g.resize(nrb); + dg.resize(nrb); + for (int k = 0; k < nrb; k++) { + g[k] = Tc * fc; + dg[k] = dTc * fc + Tc * dfc; + const double Tn = 2.0 * t * Tc - Tp, dTn = 2.0 * (2.0 * Tc + t * dTc) - dTp; + Tp = Tc; dTp = dTc; Tc = Tn; dTc = dTn; + } +} + +// the radial MLP (hidden layers silu, or tanh with act 1; linear output), as the device kernels +template +void mlp_x(const Layers &layers, std::vector h, std::vector dh, + std::vector &out, std::vector &dout, int act) +{ + const int n = layers.size(); + for (int il = 0; il < n; il++) { + const auto &L = layers[il]; + std::vector s(L.n_out, 0.0), ds(L.n_out, 0.0); + for (int k = 0; k < L.n_in; k++) + for (int j = 0; j < L.n_out; j++) { + const double w = L.W[(size_t) k * L.n_out + j]; + s[j] += w * h[k]; + ds[j] += w * dh[k]; + } + for (int j = 0; j < L.n_out; j++) { + s[j] *= L.norm; + ds[j] *= L.norm; + if (il < n - 1 && act == 1) { + s[j] = std::tanh(s[j]); + ds[j] *= 1.0 - s[j] * s[j]; + } else if (il < n - 1) { + const double sg = 1.0 / (1.0 + std::exp(-s[j])); + ds[j] = sg * (1.0 + s[j] * (1.0 - sg)) * ds[j]; + s[j] = s[j] * sg; + } + } + h.swap(s); + dh.swap(ds); + } + out.swap(h); + dout.swap(dh); +} + +} // namespace grace diff --git a/libgrace/src/fused_1l.cpp b/libgrace/src/fused_1l.cpp index b31aec1..4ee1dfc 100644 --- a/libgrace/src/fused_1l.cpp +++ b/libgrace/src/fused_1l.cpp @@ -6,6 +6,7 @@ #include "grace/bonds.h" #include "grace/fused_cpu.h" +#include "grace/radial_x.h" #include "grace/type_map.h" #include @@ -109,7 +110,7 @@ void Grace1LFused::build_splines() std::vector fv((size_t) (N + 1) * nout), dv((size_t) (N + 1) * nout), g, dg, out, dout; for (int k = 0; k <= N; k++) { radial_basis_x(k * h, m.radial_basis_p, m.nradbase, g, dg); - mlp_x(layers, g, dg, out, dout); + mlp_x(layers, g, dg, out, dout, m.radial_activation); for (int o = 0; o < nout; o++) { fv[(size_t) k * nout + o] = out[o]; dv[(size_t) k * nout + o] = h * dout[o]; @@ -162,6 +163,7 @@ void Grace1LFused::build_model() sh_coefs(lmax, sh_a, sh_b, sh_c, sh_d); build_splines(); + aaa_a = m.AAA_right_A; build_prod(pAA, m.prod_AA); build_prod(pAAA, m.prod_AAA); build_prod(pAAAA, m.prod_AAAA); @@ -205,7 +207,7 @@ void Grace1LFused::build_model() ci_of(rRho, 1, nAAv); prune(pAA, nAAv); auto liveA1 = operands(pAA, nlm, true, true); - const auto liveA1r = operands(pAAA, nlm, false, true); + const auto liveA1r = operands(pAAA, nlm, false, !aaa_a); for (int q = 0; q < nlm; q++) liveA1[q] = liveA1[q] || liveA1r[q]; build_fc(fcA1, m.fc_A1, CA, CA, PA, PA, PF, liveA1, cx); build_fc(fcAA1, m.fc_AA1, CF, CA, PF, PA, PF, liveAA1, cx); @@ -236,6 +238,7 @@ void Grace1LFused::build_model() throw Error(braces("GRACE-1L: energy MLP layer {} is {} x {} (widths 1..256, consecutive layers must match)", L, e_nin[L], e_nout[L])); e_act = m.energy_mlp_activation; + e_ln = m.energy_mlp_ln; out_scale = m.output_scale; log.info(braces("[GRACE-1L] lmax {} | A {} ch, FC {} ch | live A1 {}/{} AA {}/{} AA1 {}/{} AAA {}/{} AA2 {}/{} " "AAAA {}/{} | rho {} | spline N {}", @@ -278,17 +281,18 @@ void Grace1LFused::forward(Work &w) { PROF(T_PROD1); prod_fwd(pAA, w.A1.data(), w.A1.data(), w.AA.data(), PF); } { PROF(T_FC1); fc_fwd(fcAA1, w.AA.data(), A, w.AA1.data()); } { PROF(T_FC1); fc_fwd(fcAA2, w.AA.data(), A, w.AA2.data()); } + const float *Ar = aaa_a ? A : w.A1.data(); // AAA = AA1 x A1 (legacy: x A) if (rowred) { const float *X[4] = {A, w.AA.data(), nullptr, nullptr}; { PROF(T_RED1); red_fwd(rRho, ty, X, w.Rho.data()); } const Red *reds[1] = {&rRho}; float *outs[1] = {w.Rho.data()}; { const RowRed *rrs[1] = {&rrRho_3}; - PROF(T_PROD1); prod_red_fwd(pAAA, w.AA1.data(), w.A1.data(), PF, ty, reds, rrs, 1, outs, w.row.data()); } + PROF(T_PROD1); prod_red_fwd(pAAA, w.AA1.data(), Ar, PF, ty, reds, rrs, 1, outs, w.row.data()); } { const RowRed *rrs[1] = {&rrRho_4}; PROF(T_PROD1); prod_red_fwd(pAAAA, w.AA2.data(), w.AA2.data(), PF, ty, reds, rrs, 1, outs, w.row.data()); } } else { - { PROF(T_PROD1); prod_fwd(pAAA, w.AA1.data(), w.A1.data(), w.AAA.data(), PF); } + { PROF(T_PROD1); prod_fwd(pAAA, w.AA1.data(), Ar, w.AAA.data(), PF); } { PROF(T_PROD1); prod_fwd(pAAAA, w.AA2.data(), w.AA2.data(), w.AAAA.data(), PF); } const float *X[4] = {A, w.AA.data(), w.AAA.data(), w.AAAA.data()}; { PROF(T_RED1); red_fwd(rRho, ty, X, w.Rho.data()); } @@ -320,6 +324,8 @@ void Grace1LFused::backward(Work &w) } const Red *reds[1] = {&rRho}; const float *gouts[1] = {w.gRho.data()}; + const float *Ar = aaa_a ? w.A.data() : w.A1.data(); // AAA = AA1 x A1 (legacy: x A) + float *gAr = aaa_a ? gA : gA1; if (rowred) { const RowRed *rrs[1] = {&rrRho_4}; PROF(T_PROD1B); @@ -333,10 +339,10 @@ void Grace1LFused::backward(Work &w) if (rowred) { const RowRed *rrs[1] = {&rrRho_3}; PROF(T_PROD1B); - prod_red_bwd(pAAA, w.AA1.data(), w.A1.data(), PF, ty, reds, rrs, 1, gouts, w.gAA1.data(), gA1, w.row.data()); + prod_red_bwd(pAAA, w.AA1.data(), Ar, PF, ty, reds, rrs, 1, gouts, w.gAA1.data(), gAr, w.row.data()); } else { PROF(T_PROD1B); - prod_bwd(pAAA, w.AA1.data(), w.A1.data(), gAAA, w.gAA1.data(), gA1, PF); + prod_bwd(pAAA, w.AA1.data(), Ar, gAAA, w.gAA1.data(), gAr, PF); } { PROF(T_FC1B); fc_bwd(fcAA1, w.gAA1.data(), gAA, gA); } { PROF(T_PROD1B); prod_bwd(pAA, w.A1.data(), w.A1.data(), gAA, gA1, gA1, PF); } @@ -375,6 +381,8 @@ double Grace1LFused::energy_atom(const float *rho, int ty, float *grho) const const int nlay = eW.size(); double act[8][MAXW], pre[8][MAXW]; for (int k = 1; k < nrho; k++) act[0][k - 1] = rho[k]; + double x0[MAXW], q = 0.0; + if (e_ln) q = ln_fwd(act[0], x0, nrho - 1); for (int L = 0; L < nlay; L++) { for (int c = 0; c < e_nout[L]; c++) pre[L][c] = 0.0; for (int a = 0; a < e_nin[L]; a++) { @@ -412,6 +420,7 @@ double Grace1LFused::energy_atom(const float *rho, int ty, float *grho) const } for (int a = 0; a < e_nin[L]; a++) g[a] = gn[a]; } + if (e_ln) ln_bwd(g, x0, nrho - 1, q); grho[0] = (float) out_scale; for (int k = 1; k < nrho; k++) grho[k] = (float) g[k - 1]; return Ei; diff --git a/libgrace/src/fused_2l.cpp b/libgrace/src/fused_2l.cpp index 4088274..bdd6f55 100644 --- a/libgrace/src/fused_2l.cpp +++ b/libgrace/src/fused_2l.cpp @@ -6,6 +6,7 @@ #include "grace/bonds.h" #include "grace/fused_cpu.h" +#include "grace/radial_x.h" #include "grace/type_map.h" #include @@ -126,7 +127,7 @@ void Grace2LFused::build_splines() std::vector fv((size_t) (N + 1) * nout), dv((size_t) (N + 1) * nout), g, dg, out, dout; for (int k = 0; k <= N; k++) { radial_basis_x(k * h, m.radial_basis_p, m.nradbase, g, dg); - mlp_x(layers, g, dg, out, dout); + mlp_x(layers, g, dg, out, dout, m.radial_activation); for (int o = 0; o < nout; o++) { fv[(size_t) k * nout + o] = out[o]; dv[(size_t) k * nout + o] = h * dout[o]; @@ -359,6 +360,9 @@ void Grace2LFused::build_model() // norms + energy MLP lnI_scale = m.I_nl_LN_scale; ln0_scale = m.I_0_LN_scale; + rms = m.energy_rms_norm; + nrms1 = m.I_nl_LN_n_out - 1; + nrms0 = m.I_0_LN_n_out; if ((int) lnI_scale.size() != nrho - 1 || (int) ln0_scale.size() != nI2 || nrho != nI2) throw Error("GRACE-2L: unexpected RMSNorm widths"); eW.clear(); e_norm.clear(); e_nin.clear(); e_nout.clear(); @@ -379,6 +383,7 @@ void Grace2LFused::build_model() throw Error(braces("GRACE-2L: energy MLP layer {} is {} x {} (widths 1..256, consecutive layers must match)", L, e_nin[L], e_nout[L])); e_act = m.energy_mlp_activation; + e_ln = m.energy_mlp_ln; out_scale = m.output_scale; log.info(braces("[GRACE-2L] lmax {} | L1 {} ch, A1/AA/AAA/AAAA {}/{}/{}/{} | I {}x{} | L2 {} / {} ch, " @@ -546,11 +551,13 @@ double Grace2LFused::energy_atom(const float *rho, const float *I2, int ty, doub double ms1 = 0.0, ms0 = 0.0; for (int k = 1; k < nr; k++) ms1 += (double) rho[k] * rho[k]; for (int k = 0; k < nr; k++) ms0 += (double) I2[k] * I2[k]; - const double s1 = std::sqrt(ms1 / (nr - 1) + 1e-8), s0 = std::sqrt(ms0 / nr + 1e-8); + const double s1 = rms ? std::sqrt(ms1 / nrms1 + 1e-8) : 1.0, s0 = rms ? std::sqrt(ms0 / nrms0 + 1e-8) : 1.0; constexpr int MAXW = 256; const int nlay = eW.size(); double act[8][MAXW], pre[8][MAXW]; for (int k = 1; k < nr; k++) act[0][k - 1] = rho[k] / s1 * lnI_scale[k - 1] + I2[k] / s0 * ln0_scale[k]; + double x0[MAXW], q = 0.0; + if (e_ln) q = ln_fwd(act[0], x0, nr - 1); for (int L = 0; L < nlay; L++) { for (int c = 0; c < e_nout[L]; c++) pre[L][c] = 0.0; for (int a = 0; a < e_nin[L]; a++) { @@ -589,17 +596,20 @@ double Grace2LFused::energy_atom(const float *rho, const float *I2, int ty, doub } for (int a = 0; a < e_nin[L]; a++) g[a] = gn[a]; } - // g = dE/dx_in; RMSNorm reverses (y_k = x_k s_k / sigma) + if (e_ln) ln_bwd(g, x0, nr - 1, q); + // g = dE/dx_in; RMSNorm reverses (y_k = x_k s_k / sigma; without the RMSNorms t1 = t0 = 0, sigma = 1) grho[0] = out_scale; double t1 = 0.0; - for (int k = 1; k < nr; k++) t1 += g[k - 1] * lnI_scale[k - 1] * rho[k]; - for (int k = 1; k < nr; k++) grho[k] = g[k - 1] * lnI_scale[k - 1] / s1 - rho[k] * t1 / ((nr - 1) * s1 * s1 * s1); + if (rms) + for (int k = 1; k < nr; k++) t1 += g[k - 1] * lnI_scale[k - 1] * rho[k]; + for (int k = 1; k < nr; k++) grho[k] = g[k - 1] * lnI_scale[k - 1] / s1 - rho[k] * t1 / (nrms1 * s1 * s1 * s1); double gy[MAXW]; gy[0] = out_scale; for (int k = 1; k < nr; k++) gy[k] = g[k - 1]; double t0 = 0.0; - for (int k = 0; k < nr; k++) t0 += gy[k] * ln0_scale[k] * I2[k]; - for (int k = 0; k < nr; k++) gI2[k] = (float) (gy[k] * ln0_scale[k] / s0 - I2[k] * t0 / (nr * s0 * s0 * s0)); + if (rms) + for (int k = 0; k < nr; k++) t0 += gy[k] * ln0_scale[k] * I2[k]; + for (int k = 0; k < nr; k++) gI2[k] = (float) (gy[k] * ln0_scale[k] / s0 - I2[k] * t0 / (nrms0 * s0 * s0 * s0)); return Ei; } diff --git a/libgrace/src/model_1l.cpp b/libgrace/src/model_1l.cpp index 85c9044..33017da 100644 --- a/libgrace/src/model_1l.cpp +++ b/libgrace/src/model_1l.cpp @@ -4,6 +4,7 @@ #include "grace/error.h" #include "npz3.h" +#include "zero_pad.h" #include @@ -186,6 +187,10 @@ void GRACE1LModel::load_arrays(const cnpy::npz_t &npz) { load_mlp("energy_mlp", npz_get_int_in(npz, "energy_mlp_n_layers", 2, 1, kMaxDim), energy_mlp_layers, rho_n_out - 1, 1, "rho_n_out - 1", "one energy"); energy_mlp_activation = npz_get_int_in(npz, "energy_mlp_activation", 0, 0, 1); // 0=silu, 1=tanh + radial_activation = npz_get_int_in(npz, "radial_activation", 0, 0, 1); + AAA_right_A = npz_get_flag(npz, "AAA_right_A", 0); + if (AAA_right_A) require_eq("fc_AA1_n_out", fc_AA1.n_out, nradmax, "nradmax (AAA = AA1 x A)"); + energy_mlp_ln = fold_ln_scale(npz, energy_mlp_layers[0]); // Shifts shift_values = npz_get_double(npz, "shift_values"); @@ -274,6 +279,25 @@ void GRACE1LModel::validate_uq() const check_uq_layout("GRACE-1L", uq_layout(), uq_d_basis, uq_n_density, uq_n_elements, n_elements, {rho_instruction_names}); } +// Zero channels stay zero through every FC, CG product and reduce, so the results do not change (rho keeps its width) +GRACE1LModel GRACE1LModel::padded(int nr, int no) const +{ + if (has_uq) throw Error("GRACE-1L: only models without UQ are padded"); + if (nr < nradmax || no < fc_A1.n_out) throw Error("GRACE-1L: padding to narrower widths"); + GRACE1LModel m = *this; + const int nl = lmax + 1; + pad_radial_out(m.mlp_rad_layers.back(), nradmax, nr, nl); + m.A_lin_transform_W = zero_pad(A_lin_transform_W, {embedding_size, nradmax}, {embedding_size, nr}); + pad_fc(m.fc_A1, no, nr, nr); // A1 = FC(A, A) + pad_fc(m.fc_AA1, no, no, nr); // AA1, AA2 = FC(AA, A) + pad_fc(m.fc_AA2, no, no, nr); + pad_reduce(m.reduce_A, n_elements, rho_n_out, rho_n_out, nr); + for (ReduceWeights *w : {&m.reduce_AA, &m.reduce_AAA, &m.reduce_AAAA}) + pad_reduce(*w, n_elements, rho_n_out, rho_n_out, no); + m.nradmax = nr; + return m; +} + std::vector GRACE1LModel::uq_layout() const { std::vector L; diff --git a/libgrace/src/model_2l.cpp b/libgrace/src/model_2l.cpp index d087e33..99994a2 100644 --- a/libgrace/src/model_2l.cpp +++ b/libgrace/src/model_2l.cpp @@ -4,6 +4,7 @@ #include "grace/error.h" #include "npz3.h" +#include "zero_pad.h" #include #include @@ -186,9 +187,15 @@ void GRACE2LModel::load_arrays(const cnpy::npz_t &npz) { load_reduce("rho", "AAAA", rho_reduce_AAAA); // I_nl_LN RMSNorm - I_nl_LN_scale = npz_require_double(npz, "I_nl_LN_scale"); - I_nl_LN_type = npz_get_int_scalar(npz, "I_nl_LN_type", 1); - I_nl_LN_n_out = npz_get_int_scalar(npz, "I_nl_LN_n_out", 17); + energy_rms_norm = npz_get_flag(npz, "energy_rms_norm", 1); + if (energy_rms_norm) { + I_nl_LN_scale = npz_require_double(npz, "I_nl_LN_scale"); + I_nl_LN_type = npz_get_int_scalar(npz, "I_nl_LN_type", 1); + I_nl_LN_n_out = npz_get_int_scalar(npz, "I_nl_LN_n_out", 17); + } else { + I_nl_LN_scale.assign(rho_n_out - 1, 1.0); + I_nl_LN_n_out = rho_n_out; + } // YI coupling { @@ -218,9 +225,14 @@ void GRACE2LModel::load_arrays(const cnpy::npz_t &npz) { load_reduce("I2", "BBBB", I2_reduce_BBBB); // I_0_LN RMSNorm - I_0_LN_scale = npz_require_double(npz, "I_0_LN_scale"); - I_0_LN_type = npz_get_int_scalar(npz, "I_0_LN_type", 0); - I_0_LN_n_out = npz_get_int_scalar(npz, "I_0_LN_n_out", 17); + if (energy_rms_norm) { + I_0_LN_scale = npz_require_double(npz, "I_0_LN_scale"); + I_0_LN_type = npz_get_int_scalar(npz, "I_0_LN_type", 0); + I_0_LN_n_out = npz_get_int_scalar(npz, "I_0_LN_n_out", 17); + } else { + I_0_LN_scale.assign(I2_n_out, 1.0); + I_0_LN_n_out = I2_n_out; + } // Source names of the reduces (I2 / rho: kept for the UQ basis order) I2_instruction_names = npz_names(npz, "I2_instruction_names"); @@ -366,6 +378,8 @@ void GRACE2LModel::load_arrays(const cnpy::npz_t &npz) { // Energy MLP: rho[1:] + I2[1:] -> one energy load_mlp("energy_mlp", 2, rho_n_out - 1, 1, "rho_n_out - 1", "one energy", energy_mlp_layers); energy_mlp_activation = npz_get_int_in(npz, "energy_mlp_activation", 1, 0, 1); // 0 = silu, 1 = tanh + radial_activation = npz_get_int_in(npz, "radial_activation", 0, 0, 1); + energy_mlp_ln = fold_ln_scale(npz, energy_mlp_layers[0]); // Shifts shift_values = npz_get_double(npz, "shift_values"); @@ -457,6 +471,46 @@ void GRACE2LModel::validate_uq() const {I2_instruction_names, rho_instruction_names}); } +// Zero channels stay zero through every FC, CG product and reduce, so the results do not change. The RMSNorms of rho +// and I2 keep their unpadded counts (I_nl_LN_n_out, I_0_LN_n_out); the legacy RMS layer norm of the energy-MLP input +// divides by the input count: its folded scale (layer 0) takes sqrt(n / n') +GRACE2LModel GRACE2LModel::padded(const Widths &w) const +{ + const auto [no1, no2, i1, rho, nr2] = w; + if (has_uq) throw Error("GRACE-2L: only models without UQ are padded"); + if (no1 < L1_nradmax || no2 < B_n_out || i1 < I1_n_out || rho < rho_n_out || nr2 < L2_nradmax) + throw Error("GRACE-2L: padding to narrower widths"); + GRACE2LModel m = *this; + const int nl = lmax + 1; + pad_radial_out(m.L1_mlp_rad_layers.back(), L1_nradmax, no1, nl); + pad_radial_out(m.L2_mlp_rad_layers.back(), L2_nradmax, nr2, nl); + m.A_lin_transform_W = zero_pad(A_lin_transform_W, {embedding_size, L1_nradmax}, {embedding_size, no1}); + m.B0_lin_transform_W = zero_pad(B0_lin_transform_W, {embedding_size, L2_nradmax}, {embedding_size, nr2}); + for (FCWeights *w : {&m.fc_A1, &m.fc_AA1, &m.fc_AA2}) pad_fc(*w, no1, no1, no1); + for (FCWeights *w : {&m.fc_B1, &m.fc_BB1, &m.fc_BB2}) pad_fc(*w, no2, no2, no2); + for (ReduceWeights *w : {&m.I1_reduce_A, &m.I1_reduce_AA, &m.I1_reduce_AAA, &m.I1_reduce_AAAA}) + pad_reduce(*w, n_elements, I1_n_out, i1, no1); + for (ReduceWeights *w : {&m.rho_reduce_A, &m.rho_reduce_AA, &m.rho_reduce_AAA, &m.rho_reduce_AAAA}) + pad_reduce(*w, n_elements, rho_n_out, rho, no1); + pad_reduce(m.I_reduce_I1, 1, I_n_out, nr2, i1); + for (ReduceWeights *w : {&m.B_reduce_YI, &m.B_reduce_B0}) pad_reduce(*w, 1, B_n_out, no2, nr2); + for (ReduceWeights *w : {&m.I2_reduce_B, &m.I2_reduce_BB, &m.I2_reduce_BBB, &m.I2_reduce_BBBB}) + pad_reduce(*w, n_elements, I2_n_out, rho, no2); + m.L1_nradmax = m.L1_fc_n_out = no1; + m.L2_nradmax = m.I_n_out = m.yi.n_rad_max = nr2; + m.B_n_out = m.L2_fc_n_out = no2; + m.I1_n_out = i1; + m.rho_n_out = m.I2_n_out = rho; + m.I_nl_LN_scale = zero_pad(I_nl_LN_scale, {rho_n_out - 1}, {rho - 1}); + m.I_0_LN_scale = zero_pad(I_0_LN_scale, {I2_n_out}, {rho}); + MLPLayer &e0 = m.energy_mlp_layers[0]; // input rho[1:] + I2[1:] + e0.W = zero_pad(e0.W, {e0.n_in, e0.n_out}, {rho - 1, e0.n_out}); + if (energy_mlp_ln) + for (double &v : e0.W) v *= std::sqrt((double) e0.n_in / (rho - 1)); + e0.n_in = rho - 1; + return m; +} + std::vector GRACE2LModel::uq_layout() const { std::vector L; diff --git a/libgrace/src/model_3l.cpp b/libgrace/src/model_3l.cpp index aef1f81..b331cfc 100644 --- a/libgrace/src/model_3l.cpp +++ b/libgrace/src/model_3l.cpp @@ -166,6 +166,8 @@ void GRACE3LModel::load_arrays(const cnpy::npz_t &npz) { // a second load into the same object must keep nothing of the previous model *this = GRACE3LModel(); check_grace_format(npz, "GRACE-3L"); + for (const char *k : {"energy_mlp_scale", "radial_activation", "energy_rms_norm", "AAA_right_A"}) // format 2: 1L / 2L only + if (npz.count(k)) throw bad_array(k, "is a legacy GRACE-1L/2L key, not supported in GRACE-3L"); // ---- Global metadata ---- n_elements = get_size(npz, "n_elements"); diff --git a/libgrace/src/npz3.h b/libgrace/src/npz3.h index 8f2fd4b..1e33edf 100644 --- a/libgrace/src/npz3.h +++ b/libgrace/src/npz3.h @@ -260,7 +260,7 @@ inline int npz_get_int_in(const cnpy::npz_t &npz, const std::string &key, int de } // the newest grace_format_version of a GRACE kokkos.npz this build reads -constexpr int kGraceFormatVersion = 1; +constexpr int kGraceFormatVersion = 2; // 2: radial_activation / energy_rms_norm (legacy GRACE-1L/2L) // grace_model_kind / grace_format_version of a GRACE kokkos.npz (written by tensorpotential exports; older exports have // neither and load as before): the kind must be the loader's, the format one this build reads @@ -325,6 +325,18 @@ template inline void require_size(const std::string &key, const V &v, s throw bad_array(key, "has " + std::to_string(v.size()) + " values, expected " + std::to_string(n)); } +// legacy GRACE-1L/2L: the scale s of the RMS layer norm x s / rms(x) of the energy-MLP input (energy_mlp_scale) folded +// into the rows of its first layer L (W [n_in][n_out]); true if the model has the norm +template inline bool fold_ln_scale(const cnpy::npz_t &npz, Layer &L) +{ + const auto s = npz_get_double(npz, "energy_mlp_scale"); + if (s.empty()) return false; + require_size("energy_mlp_scale", s, L.n_in); + for (int k = 0; k < L.n_in; k++) + for (int j = 0; j < L.n_out; j++) L.W[(size_t) k * L.n_out + j] *= s[k]; + return true; +} + // got == want (why: where want comes from); noun: "has " (a dimension of the array), else "is " inline void require_eq(const std::string &key, int64_t got, int64_t want, const std::string &why, const char *noun = nullptr) diff --git a/libgrace/src/zero_pad.h b/libgrace/src/zero_pad.h new file mode 100644 index 0000000..8263b56 --- /dev/null +++ b/libgrace/src/zero_pad.h @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: LicenseRef-ASL (Academic Software Licence, see LICENSE.md) +// zero padding of weight arrays (GRACE1LModel::padded, GRACE2LModel::padded; private to libgrace) +#pragma once + +#include "grace/error.h" + +#include +#include + +namespace grace { + +// a row-major array of dims `from` zero-padded to dims `to` (to[k] >= from[k]) +inline std::vector zero_pad(const std::vector &a, const std::vector &from, + const std::vector &to) +{ + size_t n = 1, n0 = 1; + for (int t : to) n *= (size_t) t; + for (int f : from) n0 *= (size_t) f; + if (n0 != a.size()) throw Error("padding an array of unexpected size"); + std::vector out(n, 0.0); + std::vector idx(from.size(), 0); + for (double v : a) { + size_t o = 0; + for (size_t k = 0; k < to.size(); k++) o = o * to[k] + idx[k]; + out[o] = v; + for (int k = (int) from.size() - 1; k >= 0 && ++idx[k] == from[k]; k--) idx[k] = 0; + } + return out; +} + +// the last radial-MLP layer (outputs o = n * nl + l) zero-padded from n to `to` channels +template void pad_radial_out(Layer &L, int n, int to, int nl) +{ + L.W = zero_pad(L.W, {L.n_in, n, nl}, {L.n_in, to, nl}); + L.n_out = to * nl; +} + +// an FC (w_left / w_right [n_out, n_in, w_shape]) zero-padded to `no` outputs and to_left / to_right inputs +template void pad_fc(FC &w, int no, int to_left, int to_right) +{ + const int nl_in = (int) (w.w_left.size() / ((size_t) w.n_out * w.w_shape_left)); + const int nr_in = (int) (w.w_right.size() / ((size_t) w.n_out * w.w_shape_right)); + w.w_left = zero_pad(w.w_left, {w.n_out, nl_in, w.w_shape_left}, {no, to_left, w.w_shape_left}); + w.w_right = zero_pad(w.w_right, {w.n_out, nr_in, w.w_shape_right}, {no, to_right, w.w_shape_right}); + w.n_out = no; +} + +// a reduce (W [ne, n_out, n_in, w_shape]) zero-padded to to_out outputs and to_in inputs +template void pad_reduce(Reduce &w, int ne, int n_out, int to_out, int to_in) +{ + const int ws = w.w_shape > 0 ? w.w_shape : 1, n_in = (int) (w.W.size() / ((size_t) ne * n_out * ws)); + w.W = zero_pad(w.W, {ne, n_out, n_in, ws}, {ne, to_out, to_in, ws}); + w.n_in = to_in; +} + +} // namespace grace diff --git a/pyproject.toml b/pyproject.toml index 7711667..473d0f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "grace-runtime" -version = "0.1.1" +version = "0.1.2" description = "Fast CPU and GPU inference of GRACE and PACE interatomic potentials with ASE, batching and TorchSim" readme = "README.md" requires-python = ">=3.10" diff --git a/python-cuda/grace_runtime_cuda12/__init__.py b/python-cuda/grace_runtime_cuda12/__init__.py index b479a85..0a7e31a 100644 --- a/python-cuda/grace_runtime_cuda12/__init__.py +++ b/python-cuda/grace_runtime_cuda12/__init__.py @@ -4,7 +4,7 @@ import re from pathlib import Path -ABI = 6 # = grace_runtime._cuda.ABI (the C++ interface kk_api.h of the modules) +ABI = 7 # = grace_runtime._cuda.ABI (the C++ interface kk_api.h of the modules) try: from importlib.metadata import PackageNotFoundError, version diff --git a/python-cuda/pyproject.toml b/python-cuda/pyproject.toml index c391723..e15118d 100644 --- a/python-cuda/pyproject.toml +++ b/python-cuda/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "scikit_build_core.build" [project] name = "grace-runtime-cuda12" -version = "0.1.1" # = grace-runtime (scripts/build_cuda_wheels.sh sets both from the repository pyproject) +version = "0.1.2" # = grace-runtime (scripts/build_cuda_wheels.sh sets both from the repository pyproject) description = "NVIDIA GPU backend (CUDA 12) of grace-runtime: fused GRACE-1L/2L/3L Kokkos kernels for V100, RTX 20xx/30xx/40xx/50xx, A100/A30, L40S, H100/H200, B200/B300" readme = "README.md" requires-python = ">=3.10" license = { text = "Academic Software License (ASL); bundled third-party code (Kokkos: Apache-2.0 WITH LLVM-exception): see THIRD_PARTY_NOTICES.md" } classifiers = ["License :: Other/Proprietary License"] -dependencies = ["grace-runtime==0.1.1"] +dependencies = ["grace-runtime==0.1.2"] [project.urls] Homepage = "https://github.com/ICAMS/grace-runtime" diff --git a/python-rocm6/grace_runtime_rocm6/__init__.py b/python-rocm6/grace_runtime_rocm6/__init__.py index 078deb0..1577155 100644 --- a/python-rocm6/grace_runtime_rocm6/__init__.py +++ b/python-rocm6/grace_runtime_rocm6/__init__.py @@ -5,7 +5,7 @@ import re from pathlib import Path -ABI = 6 # = grace_runtime._gpu.ABI (the C++ interface kk_api.h of the modules) +ABI = 7 # = grace_runtime._gpu.ABI (the C++ interface kk_api.h of the modules) PLATFORM = "rocm" # grace_runtime._gpu.plugins: an AMD backend (a plug-in without PLATFORM is CUDA) ROCM_MAJOR = 6 # soname major of the HIP runtime the modules link (libamdhip64.so.6) diff --git a/python-rocm6/pyproject.toml b/python-rocm6/pyproject.toml index 85e5125..b299c9f 100644 --- a/python-rocm6/pyproject.toml +++ b/python-rocm6/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "scikit_build_core.build" [project] name = "grace-runtime-rocm6" -version = "0.1.1" # = grace-runtime (scripts/build_rocm_wheels.sh sets both from the repository pyproject) +version = "0.1.2" # = grace-runtime (scripts/build_rocm_wheels.sh sets both from the repository pyproject) description = "AMD GPU backend (ROCm 6) of grace-runtime: fused GRACE-1L/2L/3L Kokkos kernels for MI210/MI250/MI250X and MI300A/MI300X" readme = "README.md" requires-python = ">=3.10" license = { text = "Academic Software License (ASL); bundled third-party code (Kokkos: Apache-2.0 WITH LLVM-exception): see THIRD_PARTY_NOTICES.md" } classifiers = ["License :: Other/Proprietary License"] -dependencies = ["grace-runtime==0.1.1"] +dependencies = ["grace-runtime==0.1.2"] [project.urls] Homepage = "https://github.com/ICAMS/grace-runtime" diff --git a/python/grace_runtime/_gpu.py b/python/grace_runtime/_gpu.py index aa67049..ba149f9 100644 --- a/python/grace_runtime/_gpu.py +++ b/python/grace_runtime/_gpu.py @@ -15,7 +15,7 @@ import sys import warnings -ABI = 6 # the C++ interface of the GPU modules (kk_api.h); a plug-in must carry the same +ABI = 7 # the C++ interface of the GPU modules (kk_api.h); a plug-in must carry the same # the GPU module in use: platform "cuda" | "rocm", device index, module arch, ptx (CUDA only), cc (CUDA: "8.6"; ROCm: # the gfx target), driver (CUDA driver version, set by _cuda.compute_capability) _state = {"module": None, "platform": None, "device": None, "arch": None, "ptx": False, "cc": None, "driver": None} @@ -32,13 +32,18 @@ def strict() -> bool: return _flag("GRACE_RUNTIME_STRICT") -def notice(msg: str, stacklevel: int = 3): +def notice(msg: str, stacklevel: int | None = 3): """A fallback the user must not miss: a banner on stderr once per process and message, not subject to warning - filters (GRACE_RUNTIME_QUIET=1 silences it), then a RuntimeWarning (for warning filters and pytest.warns).""" + filters (GRACE_RUNTIME_QUIET=1 silences it), then a RuntimeWarning (for warning filters and pytest.warns). + stacklevel None: the warning points at the first caller outside grace_runtime.""" if msg not in _noticed and not _flag("GRACE_RUNTIME_QUIET"): _noticed.add(msg) bar = "!" * 100 print(f"{bar}\ngrace-runtime WARNING: {msg}\n{bar}", file=sys.stderr, flush=True) + if stacklevel is None: + pkg, f, stacklevel = os.path.dirname(os.path.abspath(__file__)) + os.sep, sys._getframe(), 1 + while f.f_back is not None and os.path.abspath(f.f_code.co_filename).startswith(pkg): + f, stacklevel = f.f_back, stacklevel + 1 warnings.warn(msg, RuntimeWarning, stacklevel=stacklevel) @@ -168,17 +173,18 @@ def info() -> dict | None: "native": not _state["ptx"]} -_unsupported = {} # model path -> why the fused GPU kernels refuse it +_unsupported = {} # (model path, GRACE_GENERIC_KERNELS) -> why the fused GPU kernels refuse it def unsupported(mod, kk_model, path, spline): """None when the fused GPU kernels take the model, else the reason; probed once per path (the engine of the first probe's spline is set up then, so it costs nothing later)""" - if path not in _unsupported: + key = (path, os.environ.get("GRACE_GENERIC_KERNELS")) + if key not in _unsupported: try: mod.System(kk_model, [kk_model.element_names[0]], spline=spline) - _unsupported[path] = None + _unsupported[key] = None except ValueError as e: - _unsupported[path] = str(e) - return _unsupported[path] + _unsupported[key] = str(e) + return _unsupported[key] diff --git a/python/grace_runtime/calculator.py b/python/grace_runtime/calculator.py index cb0c5b7..9b2a642 100644 --- a/python/grace_runtime/calculator.py +++ b/python/grace_runtime/calculator.py @@ -77,14 +77,18 @@ def _kk_model_key(model) -> str: def _load_kk_model(kk, model): asi = getattr(model, "asi", None) - return kk.Model(model.path, asi=asi) if asi else kk.Model(model.path) + m = kk.Model(model.path, asi=asi) if asi else kk.Model(model.path) + if note := m.slow_notice: + _gpu.notice(note, stacklevel=None) # from the calculator or the batch functions + return m def _select_device(device: str, model): """(resolved device, (GPU module, its model) or None, backend info) for a CPU Model. device="auto" falls back to the CPU (announced by _gpu.notice; an error with GRACE_RUNTIME_STRICT=1) when a GPU is there but has no module, the - module predates the GRACE-3L engine, the fused kernels refuse the model or the GPU fails to start / set up the model - (e.g. its memory held by another process; not remembered); an explicit GPU device ("cuda", "rocm", "gpu") raises. + module predates the GRACE-3L engine, the fused kernels refuse the model (not one the generic GPU kernels can run, + GRACE_GENERIC_KERNELS: that raises) or the GPU fails to start / set up the model (e.g. its memory held by another + process; not remembered); an explicit GPU device ("cuda", "rocm", "gpu") raises. Backend info: {"device", "platform", "arch", "cc", "native", "fallback", "store"} (fallback: None, "ptx" or the reason of a CPU fallback; store: see _store_info, set by the calculator's CPU calls).""" path = model.path @@ -112,6 +116,8 @@ def _select_device(device: str, model): kk_model = _KK_MODELS[mkey] if auto: why = _gpu.unsupported(kk, kk_model, path, model.spline) + if why and kk_model.needs_generic: # the generic GPU kernels can run it (GRACE_GENERIC_KERNELS=1) + raise ValueError(why) if why: return _cpu_fallback(why) except RuntimeError as e: diff --git a/python/grace_runtime/models.json b/python/grace_runtime/models.json index c365e83..703f2f6 100644 --- a/python/grace_runtime/models.json +++ b/python/grace_runtime/models.json @@ -2,6 +2,24 @@ "schema": 1, "generated_by": "tensorpotential", "models": [ + { + "name": "GRACE-1L-OAM", + "aliases": [ + "GRACE-1L-OAM_2Feb25" + ], + "arch": "1L", + "url": "https://huggingface.co/AMS-ICAMS-RUB/grace-foundation-models/resolve/kk-legacy-v1/kokkos/GRACE-1L-OAM-kokkos.npz", + "sha256": "81f3c70fb17a9533194069f253879ceb56730ee142e249473b0dbbb15b3c037f", + "uq": false + }, + { + "name": "GRACE-1L-OMAT", + "aliases": [], + "arch": "1L", + "url": "https://huggingface.co/AMS-ICAMS-RUB/grace-foundation-models/resolve/kk-legacy-v1/kokkos/GRACE-1L-OMAT-kokkos.npz", + "sha256": "1255e74d9851cb4a8543b58e75499d9b5d0b99236ad9f65abfcebad1b7c8950b", + "uq": false + }, { "name": "GRACE-1L-OMAT-large-base", "aliases": [ @@ -86,6 +104,24 @@ "sha256": "97a7f56dec6a00bfc042f463228efd2c15baa020b66f41a4e5e67fcd826c850c", "uq": true }, + { + "name": "GRACE-2L-OAM", + "aliases": [ + "GRACE_2L_OAM_28Jan25" + ], + "arch": "2L", + "url": "https://huggingface.co/AMS-ICAMS-RUB/grace-foundation-models/resolve/kk-legacy-v1/kokkos/GRACE-2L-OAM-kokkos.npz", + "sha256": "86ea29271df85f8955c89d6f376e335781b5423957617b550b1a5ca38e69c261", + "uq": false + }, + { + "name": "GRACE-2L-OMAT", + "aliases": [], + "arch": "2L", + "url": "https://huggingface.co/AMS-ICAMS-RUB/grace-foundation-models/resolve/kk-legacy-v1/kokkos/GRACE-2L-OMAT-kokkos.npz", + "sha256": "37501f76d03d61293a43a3c4a756977271391ad54d58105aecb7662e0ae051e9", + "uq": false + }, { "name": "GRACE-2L-OMAT-large-base", "aliases": [ @@ -208,13 +244,7 @@ } ], "unsupported": { - "GRACE-1L-OAM": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE-1L-OAM_2Feb25": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE-1L-OMAT": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE-2L-OAM": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE-2L-OMAT": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", "GRACE-FS-OAM": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE-FS-OMAT": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator", - "GRACE_2L_OAM_28Jan25": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator" + "GRACE-FS-OMAT": "no kokkos.npz export (legacy model); use tensorpotential's TPCalculator" } } diff --git a/python/kk/kk_api.h b/python/kk/kk_api.h index aaf0d9a..5aa60b9 100644 --- a/python/kk/kk_api.h +++ b/python/kk/kk_api.h @@ -54,6 +54,8 @@ class Model { ~Model(); int layers() const; // 1, 2 or 3 (GRACE); 0: a PACE model bool has_uq() const; + std::string slow_notice() const; // non-empty: the GPU runs the model zero-padded or on the generic kernels + bool needs_generic() const; // GRACE-2L outside the fused kernels: GPU only with GRACE_GENERIC_KERNELS=1 void validate_uq() const; // throws std::invalid_argument unless the UQ head is usable std::vector element_names() const; int64_t engine_setups() const; // full GPU engine setups so far (one per spline / chunksize, any species) diff --git a/python/kk/kk_core.cpp b/python/kk/kk_core.cpp index 5c484c8..89c427d 100644 --- a/python/kk/kk_core.cpp +++ b/python/kk/kk_core.cpp @@ -45,7 +45,8 @@ namespace grace_runtime_kk { using Dev = Kokkos::DefaultExecutionSpace; using T = grace::kk::DefaultTraits; using E1 = grace::kk::Grace1LFusedKK; -using E2 = grace::kk::Grace2LFusedKK; +using E2F = grace::kk::Grace2LFusedKK; +using E2 = grace::kk::Grace2LKK; // E2F, or the generic engine (E2F::use_generic) using E3 = grace::kk::Grace3LFusedKK; // as LAMMPS grace/runtime/kk (fp32 with fp64 bonds / envelope) using EP = grace::kk::PaceKK; // PACE (fp64) } // namespace grace_runtime_kk @@ -222,7 +223,9 @@ struct SharedEngine { struct Model::Impl : DeviceState { std::string layers; std::shared_ptr m1; + std::shared_ptr m1run; // m1 zero-padded to compiled widths, once (E1::padded_model) std::shared_ptr m2; + std::shared_ptr m2run; // m2 zero-padded to a compiled width set, once (E2F::padded_model) std::shared_ptr m3; std::shared_ptr mp; std::mutex emu; @@ -279,7 +282,7 @@ struct Model::Impl : DeviceState { e->e1 = std::make_unique(); e->e1->spline_n = spline; e->e1->chunksize = chunksize; - e->e1->set_model(m1, type2mu); + e->e1->set_model(m1run, type2mu); e->e1->setup(); } else if (m3) { e->e3 = std::make_unique(); @@ -288,10 +291,15 @@ struct Model::Impl : DeviceState { e->e3->set_model(m3, type2mu); e->e3->setup(); } else { - e->e2 = std::make_unique(); - e->e2->spline_n = spline; + if (E2F::use_generic(*m2)) { + e->e2 = std::make_unique(); + } else { + auto f = std::make_unique(); + f->spline_n = spline; + e->e2 = std::move(f); + } e->e2->chunksize = chunksize; - e->e2->set_model(m2, type2mu); + e->e2->set_model(m2run, type2mu); e->e2->setup(); } } catch (const grace::Error &err) { @@ -320,12 +328,14 @@ Model::Model(const std::string &path, int /*spline*/, const std::string &asi) : } else if (impl->layers == "1l") { impl->m1 = std::make_shared(); impl->m1->load(path); + impl->m1run = E1::padded_model(impl->m1); } else if (impl->layers == "3l") { impl->m3 = std::make_shared(); impl->m3->load(path); } else { impl->m2 = std::make_shared(); impl->m2->load(path); + impl->m2run = E2F::padded_model(impl->m2); } } catch (const std::exception &e) { throw std::invalid_argument(std::string("cannot load ") + path + ": " + e.what()); @@ -337,6 +347,11 @@ bool Model::has_uq() const { return impl->mp ? impl->mp->has_uq : impl->m1 ? impl->m1->has_uq : impl->m3 ? impl->m3->has_uq : impl->m2->has_uq; } +std::string Model::slow_notice() const +{ + return impl->m2 ? E2F::slow_notice(*impl->m2) : impl->m1 ? E1::slow_notice(*impl->m1) : ""; +} +bool Model::needs_generic() const { return impl->m2 && E2F::needs_generic(*impl->m2); } void Model::validate_uq() const { try { diff --git a/python/kk/kk_module.cpp b/python/kk/kk_module.cpp index 66b049a..c0208d4 100644 --- a/python/kk/kk_module.cpp +++ b/python/kk/kk_module.cpp @@ -226,6 +226,8 @@ NB_MODULE(GRACE_KK_MODULE, m) .def(nb::init(), "path"_a, "spline"_a = 1024, "asi"_a = "") .def_prop_ro("layers", &gk::Model::layers) .def_prop_ro("has_uq", &gk::Model::has_uq) + .def_prop_ro("slow_notice", &gk::Model::slow_notice) + .def_prop_ro("needs_generic", &gk::Model::needs_generic) .def_prop_ro("element_names", &gk::Model::element_names) .def_prop_ro("engine_setups", &gk::Model::engine_setups) .def("validate_uq", &gk::Model::validate_uq); diff --git a/python/tests/test_gpu_parity.py b/python/tests/test_gpu_parity.py index 74d4931..c5cca52 100644 --- a/python/tests/test_gpu_parity.py +++ b/python/tests/test_gpu_parity.py @@ -3,8 +3,10 @@ tests/fixtures/ref_gpu/_<1l|2l>_fused.{dump,thermo} from LAMMPS runs of the fused kk styles. fp32 parity: energy 2e-6 eV/atom, per-atom energy 1e-5 eV, forces 2e-5 eV/A, pressure 1e-5 relative.""" import os +import shutil import subprocess import sys +import warnings from pathlib import Path import numpy as np @@ -67,11 +69,13 @@ def test_against_lammps_fused_kk(mod, layer, cell): TF_REF = FX / "ref_tf" / "registry_tf_fp64.npz" # TPCalculator, -fp64 SavedModels; CPU side: test_ref_tf.py -@pytest.mark.parametrize("model", ["GRACE-2L-SMAX-OMAT-large"]) # its own 2L shape on the GPU (g2l_shape_smaxl.h) +# GRACE-2L-SMAX-OMAT-large: its own 2L shape on the GPU (g2l_shape_smaxl.h); the legacy models: the legacy layout (the +# fused engines; 2L at its own compiled widths) +@pytest.mark.parametrize("model", ["GRACE-2L-SMAX-OMAT-large", "GRACE-1L-OAM", "GRACE-1L-OMAT", "GRACE-2L-OAM", "GRACE-2L-OMAT"]) @pytest.mark.parametrize("cell", ["W54", "WMo54", "W2tri"]) def test_against_tf_fp64(mod, model, cell): """independent references: TPCalculator on the -fp64 SavedModel""" - path = Path(os.environ.get("GRACE_TEST_M2L") or CACHE / model / "kokkos.npz") + path = Path(model == "GRACE-2L-SMAX-OMAT-large" and os.environ.get("GRACE_TEST_M2L") or CACHE / model / "kokkos.npz") if not path.exists(): pytest.skip(f"registry model {model} not in the test cache") types, pos, h = load_cell(cell) @@ -91,6 +95,44 @@ def test_against_tf_fp64(mod, model, cell): assert np.abs(stress - ref[f"{k}/stress"]).max() / np.abs(ref[f"{k}/stress"]).max() <= 1e-5 +@pytest.mark.parametrize("name, path_kind", [("legacy_2l_small", "zero-padded"), ("legacy_2l_wide", "generic"), + ("grace_2l_small", "zero-padded"), ("legacy_1l_wide", "zero-padded"), + ("grace_1l_small", "zero-padded")]) +def test_outside_compiled_widths(mod, name, path_kind, monkeypatch, tmp_path): + """GRACE-1L/2L models with widths outside the compiled kernels (tests/fixtures/{name}.npz, random weights): + legacy_2l_small, grace_2l_small (the RMSNorms of the current models), legacy_1l_wide (42 radial channels) and + grace_1l_small (radial 20, FC 40) run zero-padded to compiled widths; legacy_2l_wide, wider than every set, is + refused (device="auto" too: no CPU fallback) unless GRACE_GENERIC_KERNELS=1, then runs on the generic GPU engine. + Either way a warning (slow_notice) says so, and the results are the CPU engine's.""" + pytest.importorskip("grace_runtime._core") + from grace_runtime import GRACERuntimeCalculator + from ase.io import read + path = str(FX / f"{name}.npz") + at = read(FX / "WMo54.data", format="lammps-data", atom_style="atomic", Z_of_type={1: 74, 2: 42}) + monkeypatch.delenv("GRACE_GENERIC_KERNELS", raising=False) + note = "zero-padded" if path_kind == "zero-padded" else "generic GPU kernels" + assert mod.Model(path).needs_generic == (path_kind == "generic") + if path_kind == "generic": + assert mod.Model(path).slow_notice == "" + monkeypatch.setenv("GRACE_RUNTIME_DEVICE", GPU) # device="auto" picks the GPU (conftest pins the CPU) + for dev in (GPU, "auto"): + with pytest.raises(ValueError, match="GRACE_GENERIC_KERNELS=1"): + at.calc = GRACERuntimeCalculator(path, device=dev) + at.get_potential_energy() + monkeypatch.setenv("GRACE_GENERIC_KERNELS", "1") + path = shutil.copy(path, tmp_path) # another path: the calculator caches the GPU model (and its notice) + assert note in mod.Model(path).slow_notice + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + at.calc = GRACERuntimeCalculator(path, device=GPU) + assert any(note in str(w.message) for w in caught) + e, f, ei = at.get_potential_energy(), at.get_forces(), at.get_potential_energies() + at.calc = GRACERuntimeCalculator(path, device="cpu", num_threads=4) + assert abs(e - at.get_potential_energy()) / len(at) <= 2e-6 + assert np.abs(f - at.get_forces()).max() <= 2e-5 + assert np.abs(ei - at.get_potential_energies()).max() <= 1e-5 + + def test_permuted_input_same_result(mod): """results come back in the caller's atom order whatever the list's spatial order""" types, pos, h = load_cell("WMo54") diff --git a/python/tests/test_pace.py b/python/tests/test_pace.py index 948d6e2..191bba5 100644 --- a/python/tests/test_pace.py +++ b/python/tests/test_pace.py @@ -137,8 +137,9 @@ def test_yaml_and_npz_same(model_files, tmp_path): b = at.copy() b.calc = GRACERuntimeCalculator(p, device="cpu", num_threads=2) r.append((b.get_potential_energy(), b.get_forces())) - assert r[0][0] == r[1][0] - np.testing.assert_array_equal(r[0][1], r[1][1]) + # two OpenMP threads sum in a run-dependent order: equal to round-off, not bit for bit + np.testing.assert_allclose(r[0][0], r[1][0], atol=1e-9, rtol=0) + np.testing.assert_allclose(r[0][1], r[1][1], atol=1e-9, rtol=0) def test_yaml_converted_on_the_fly_without_cache(tmp_path, monkeypatch): diff --git a/python/tests/test_ref_tf.py b/python/tests/test_ref_tf.py index 75a741f..e02a344 100644 --- a/python/tests/test_ref_tf.py +++ b/python/tests/test_ref_tf.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: LicenseRef-ASL (Academic Software Licence, see LICENSE.md) -"""Registry models against independent references: tensorpotential's TPCalculator on their -fp64 SavedModels, on the -fixture cells (tests/fixtures/ref_tf/registry_tf_fp64.npz). Covers the -GRACE-2L-SMAX-OMAT-large architecture, which the LAMMPS references (conftest TEST_MODELS) do not; the GPU side is -python/tests/test_gpu_parity.py::test_against_tf_fp64. Thresholds: fp32 network against fp64 (test_vs_tp.py).""" +"""Registry models against independent references: tensorpotential's TPCalculator on their -fp64 SavedModels (the +legacy GRACE-1L/2L models: their fp64 SavedModels), on the fixture cells (tests/fixtures/ref_tf/registry_tf_fp64.npz). +Covers the GRACE-2L-SMAX-OMAT-large architecture, which the LAMMPS references (conftest TEST_MODELS) do not, and the +legacy layout (tanh radial MLPs, RMS layer norm of the energy-MLP input, 1L AAA = AA1 x A, 2L without RMSNorms); the +GPU side is python/tests/test_gpu_parity.py::test_against_tf_fp64. Thresholds: fp32 network against fp64 +(test_vs_tp.py).""" import os import numpy as np @@ -13,7 +15,8 @@ FX = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "tests", "fixtures") REF = os.path.join(FX, "ref_tf", "registry_tf_fp64.npz") -MODELS = {"GRACE-2L-SMAX-OMAT-large": "GRACE_TEST_M2L"} +MODELS = ["GRACE-2L-SMAX-OMAT-large", "GRACE-1L-OAM", "GRACE-1L-OMAT", "GRACE-2L-OAM", "GRACE-2L-OMAT"] +ENV = {"GRACE-2L-SMAX-OMAT-large": "GRACE_TEST_M2L"} # name: the environment variable of another file CELLS = ["W54", "WMo54", "W2tri"] TOL = {"energy": 1e-6, "forces": 2e-5, "energies": 1e-5, "stress": 1e-5} # eV/atom, eV/A, eV, relative TOL_SMALL = 2e-6 # eV/atom below 10 atoms (fp32 floor of the total energy, as test_vs_tp.py) @@ -24,10 +27,16 @@ def ref(): return np.load(REF) -@pytest.fixture(scope="module", params=list(MODELS)) +@pytest.fixture(scope="module", params=MODELS) def model(request): from grace_runtime import registry - return request.param, os.environ.get(MODELS[request.param]) or registry.resolve(request.param) + name = request.param + try: + return name, os.environ.get(ENV.get(name, "")) or registry.resolve(name) + except ValueError as e: # no export in the registry yet (a user registry in the cache may have one) + if name not in registry.load_entries().get("unsupported", {}): + raise + pytest.skip(str(e)) @pytest.mark.parametrize("cell", CELLS) diff --git a/scripts/test_cuda_wheels.sh b/scripts/test_cuda_wheels.sh index f11f7c0..87f4693 100755 --- a/scripts/test_cuda_wheels.sh +++ b/scripts/test_cuda_wheels.sh @@ -8,7 +8,8 @@ # path); WITH_CUBLAS=0 none (steps 1 and 2 set GRACE_RUNTIME_NO_CUBLAS=1, since the host may have a toolkit); # 2. the GPU test files (parity with the CPU calculator, batches, auto device, plug-in logic, # memory, allocation failures, chunks, several calculators; TorchSim and device batches with WITH_TORCH=1, which -# installs the [torchsim] extra), +# installs the [torchsim] extra, with torch from the PyTorch CUDA 12 index on a CUDA 12 driver, and requires +# that torch sees the GPU), # a release gate: only the skips GRACE_TEST_ALLOW_SKIP matches pass (card sizes, timings, the PTX path without an # older module, the unsupported-model check without that model), so a missing model or package is a failure; # 3. an architecture without a module (GRACE_RUNTIME_CUDA_ARCH=sm61): auto -> CPU with a warning, cuda -> ImportError; @@ -47,7 +48,19 @@ check $? "install wheels" if [ "$WITH_CUBLAS" = 1 ]; then "$P" -m pip install -q "nvidia-cublas-cu12>=12.4,<13"; check $? "install nvidia-cublas-cu12" fi -if [ "${WITH_TORCH:-0}" = 1 ]; then "$P" -m pip install -q "$CPU_WHL[torchsim]"; check $? "install [torchsim]"; fi +if [ "${WITH_TORCH:-0}" = 1 ]; then + # PyPI's newest torch is built for CUDA 13 and sees no GPU on a CUDA 12 driver: torch first, from the matching index + CU=$(nvidia-smi | sed -n 's/.*CUDA Version: *\([0-9]*\).*/\1/p' | head -1) + if [ "${CU:-0}" -lt 13 ]; then + "$P" -m pip install -q torch --index-url https://download.pytorch.org/whl/cu129 + check $? "install torch (CUDA 12 driver: cu129 index)" + fi + "$P" -m pip install -q "$CPU_WHL[torchsim]"; check $? "install [torchsim]" + # numpy first, as in the tests: the env's libstdc++ (conda-forge numpy) before torch can load an older system one + "$P" -c "import numpy, torch; print('torch', torch.__version__, 'CUDA', torch.version.cuda) +assert torch.cuda.is_available(), 'torch sees no GPU (driver CUDA $CU)'" + check $? "torch sees the GPU" +fi W=$(mktemp -d /tmp/gn-cuda-test.XXXX) mkdir -p "$W/python" "$W/tests" cp -r "$ROOT/python/tests" "$W/python/" diff --git a/tests/fixtures/grace_1l_small.npz b/tests/fixtures/grace_1l_small.npz new file mode 100644 index 0000000..9a1c9ad Binary files /dev/null and b/tests/fixtures/grace_1l_small.npz differ diff --git a/tests/fixtures/grace_2l_small.npz b/tests/fixtures/grace_2l_small.npz new file mode 100644 index 0000000..637aa50 Binary files /dev/null and b/tests/fixtures/grace_2l_small.npz differ diff --git a/tests/fixtures/legacy_1l_wide.npz b/tests/fixtures/legacy_1l_wide.npz new file mode 100644 index 0000000..6654c7a Binary files /dev/null and b/tests/fixtures/legacy_1l_wide.npz differ diff --git a/tests/fixtures/legacy_2l_small.npz b/tests/fixtures/legacy_2l_small.npz new file mode 100644 index 0000000..65085ad Binary files /dev/null and b/tests/fixtures/legacy_2l_small.npz differ diff --git a/tests/fixtures/legacy_2l_wide.npz b/tests/fixtures/legacy_2l_wide.npz new file mode 100644 index 0000000..10322c0 Binary files /dev/null and b/tests/fixtures/legacy_2l_wide.npz differ diff --git a/tests/fixtures/ref_tf/registry_tf_fp64.npz b/tests/fixtures/ref_tf/registry_tf_fp64.npz index 021fdb1..04de9d0 100644 Binary files a/tests/fixtures/ref_tf/registry_tf_fp64.npz and b/tests/fixtures/ref_tf/registry_tf_fp64.npz differ diff --git a/tests/test_driver.cpp b/tests/test_driver.cpp index 716e150..b0eeb47 100644 --- a/tests/test_driver.cpp +++ b/tests/test_driver.cpp @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include namespace { @@ -272,6 +274,42 @@ int main() std::printf("3L fc_fold vs unfolded: dE %.3e max|dF| %.3e\n", std::fabs(a.e - b.e), df); CHECK(std::fabs(a.e - b.e) / c.n <= 1e-6 && df <= 1e-5); } + // test_2l_padded / test_1l_padded: a model zero-padded to wider channels (GRACE2LModel / GRACE1LModel::padded, how + // the GPU engines run it on compiled widths) gives the original's results. tests/fixtures: legacy_2l_small (radial + // channels 16 / 8), grace_2l_small (16 / 8, the RMSNorms of the current models), legacy_1l_wide (radial and FC + // 42 -> 64), grace_1l_small (radial 20 -> 32, FC 40 -> 64) + { + const Cell c = load_cell(fx + "/WMo54.npz"); + auto same = [&](auto &a, auto &b, const std::string &name) { + const Result ra = run(a, c, 1, true), rb = run(b, c, 1, true); + double df = 0, dpe = 0; + for (int q = 0; q < 3 * c.n; q++) df = std::max(df, std::fabs(ra.f[q] - rb.f[q])); + for (int i = 0; i < c.n; i++) dpe = std::max(dpe, std::fabs(ra.eatom[i] - rb.eatom[i])); + std::printf("padded vs original (%s): dE %.3e max|de_i| %.3e max|dF| %.3e\n", name.c_str(), + std::fabs(ra.e - rb.e), dpe, df); + CHECK(std::fabs(ra.e - rb.e) / c.n <= 1e-9 && dpe <= 1e-9 && df <= 1e-6); + }; + for (const auto &[name, w] : std::vector>{ + {"legacy_2l_small", {32, 32, 16, 13, 32}}, {"grace_2l_small", {42, 64, 12, 17, 32}}}) { + auto ms = std::make_shared(); + ms->load(fx + "/" + name + ".npz"); + auto mp = std::make_shared(ms->padded(w)); + grace::Grace2LFused a(ms, {}, {}), b(mp, {}, {}); + a.setup_types(2, type2mu_for(ms->element_names)); + b.setup_types(2, type2mu_for(mp->element_names)); + same(a, b, name); + } + for (const auto &[name, nr, no] : std::vector>{{"legacy_1l_wide", 64, 64}, + {"grace_1l_small", 32, 64}}) { + auto ms = std::make_shared(); + ms->load(fx + "/" + name + ".npz"); + auto mp = std::make_shared(ms->padded(nr, no)); + grace::Grace1LFused a(ms, {}, {}), b(mp, {}, {}); + a.setup_types(2, type2mu_for(ms->element_names)); + b.setup_types(2, type2mu_for(mp->element_names)); + same(a, b, name); + } + } // thread count changes the summation order only (document the size, require it to be tiny) { const Cell c = load_cell(fx + "/WMo54.npz"); diff --git a/tests/test_model_load.cpp b/tests/test_model_load.cpp index 9c927f6..53effb3 100644 --- a/tests/test_model_load.cpp +++ b/tests/test_model_load.cpp @@ -116,7 +116,9 @@ int main() grace::GRACE1LModel ok; ok.load_arrays(npz); CHECK(ok.n_elements == m1.n_elements); - set("GRACE-1L", 2); + set("GRACE-1L", 2); // format 2: legacy GRACE-1L/2L keys (absent here: the defaults) + ok.load_arrays(npz); + set("GRACE-1L", 3); CHECK(refused("update grace-runtime")); set("GRACE-2L", 1); CHECK(refused("holds a GRACE-2L model"));