diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py index e2af1595e3..fd2829b6e3 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/attention.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/attention.py @@ -28,6 +28,17 @@ ) +def _stop_gradient(value: Any) -> Any: + """Return ``value`` with backend gradient tracking disabled.""" + if array_api_compat.is_torch_array(value): + return value.detach() + if array_api_compat.is_jax_array(value): + import jax + + return jax.lax.stop_gradient(value) + return value + + def segment_envelope_gated_softmax( logits: Any, edge_env: Any, @@ -60,22 +71,21 @@ def segment_envelope_gated_softmax( Unconstrained denominator bias with shape (F, H). Softplus is applied to keep the bias strictly positive. eps - Small epsilon for denominator stability. + Small positive floor added to the physical null mass. src_weight Optional per-edge source-side multiplier with shape (E, 1) or - (E,). When provided the per-edge weight becomes - ``edge_env**2 * src_weight`` and the attention reduces to - ``edge_env**2 * src_weight * exp(logits) / - (zeta + sum(edge_env**2 * src_weight * exp(logits)))``. + (E,). When provided, the physical per-edge mass is + ``edge_env**2 * src_weight * exp(logits)`` and the denominator is the + sum of edge masses plus the positive null mass + ``softplus(z_bias_raw) + eps``. ``src_weight = 0`` therefore removes the source from both the numerator and the denominator, which is what SFPG needs so that a muted source does not even leak through the softmax normalization. edge_mask - Optional padded-edge validity mask with shape (E,) or (E, 1); - zero marks invalid slots. Folded into the non-negative per-edge - weight so invalid slots drop out of the group max, the numerator, - and the denominator. + Optional binary padded-edge validity mask with shape (E,) or (E, 1); + one marks valid slots and zero marks invalid slots. Invalid slots drop + out of the group max, the numerator, and the denominator. Returns ------- @@ -85,72 +95,82 @@ def segment_envelope_gated_softmax( xp = array_api_compat.array_namespace(logits) n_edge, n_focus, n_head = logits.shape n_channel = n_focus * n_head - eps_f = float(eps) device = array_api_compat.device(logits) + input_dtype = logits.dtype + promote = "float16" in str(input_dtype) + compute_dtype = xp.float32 if promote else input_dtype dst = xp.astype(dst, xp.int64) - # === Step 1. Flatten (F, H) and build the effective per-edge weight === - logits_2d = xp.reshape(logits, (n_edge, n_channel)) - zeros_e = xp.zeros((n_edge,), dtype=logits.dtype, device=device) - edge_env_1d = xp.astype(xp.reshape(edge_env, (n_edge,)), logits.dtype) - edge_env_1d = xp.where(edge_env_1d > 0.0, edge_env_1d, zeros_e) - # edge_weight_sq acts as the non-negative multiplier applied to every - # ``exp(logit)`` term. Folding ``src_weight`` (and, in the padded - # layout, ``edge_mask``) here guarantees that any edge with zero weight - # is excluded from the group max, the numerator, and the denominator in - # a single pass. - edge_weight_sq = edge_env_1d * edge_env_1d + # === Step 1. Build factor-wise effective logits === + # Computing the logarithms before multiplying the factors avoids losing a + # physically nonzero edge when ``edge_env**2 * src_weight`` underflows. + logits_2d = xp.astype(xp.reshape(logits, (n_edge, n_channel)), compute_dtype) + edge_env_1d = xp.astype(xp.reshape(edge_env, (n_edge,)), compute_dtype) + edge_positive = edge_env_1d > 0.0 + ones = xp.ones((n_edge,), dtype=compute_dtype, device=device) + log_weight = 2.0 * xp.log(xp.where(edge_positive, edge_env_1d, ones)) + active = edge_positive + source_ratio = None if src_weight is not None: - src_weight_1d = xp.astype(xp.reshape(src_weight, (n_edge,)), logits.dtype) - src_weight_1d = xp.where(src_weight_1d > 0.0, src_weight_1d, zeros_e) - edge_weight_sq = edge_weight_sq * src_weight_1d + source_weight = xp.astype(xp.reshape(src_weight, (n_edge,)), compute_dtype) + source_positive = source_weight > 0.0 + safe_source = xp.where(source_positive, source_weight, ones) + source_scale = _stop_gradient(safe_source) + log_weight = log_weight + xp.log(source_scale) + source_ratio = xp.where( + source_positive, + source_weight / source_scale, + xp.zeros((n_edge,), dtype=compute_dtype, device=device), + ) + active = active & source_positive if edge_mask is not None: - mask_1d = xp.astype(xp.reshape(edge_mask, (n_edge,)), logits.dtype) - edge_weight_sq = edge_weight_sq * mask_1d - zeta = xp.astype(xp.reshape(softplus_t(z_bias_raw), (1, n_channel)), logits.dtype) - has_weight = edge_weight_sq > 0.0 + mask = xp.astype(xp.reshape(edge_mask, (n_edge,)), compute_dtype) + mask_positive = mask > 0.0 + # ``edge_mask`` is a binary validity mask, so its positive branch has + # log-factor zero and only needs to participate in the active predicate. + active = active & mask_positive + + effective_logits = logits_2d + log_weight[:, None] minus_inf = xp.full( (n_edge, n_channel), float("-inf"), - dtype=logits.dtype, + dtype=compute_dtype, device=device, ) - logits_for_max = xp.where( - has_weight[:, None], - logits_2d, + effective_logits = xp.where( + active[:, None], + effective_logits, minus_inf, ) + null_mass = xp.reshape( + softplus_t(xp.astype(z_bias_raw, compute_dtype)) + float(eps), + (1, n_channel), + ) + null_logit = xp.log(null_mass) - # === Step 2. Destination-wise max for stable exponentials === - # Destination segment max over ``dst`` (pt ``scatter_reduce`` amax). The - # scatter is layout-agnostic and the maximum is order-independent, so the - # padded ``call`` stays bit-exact while the sparse ``call_with_edges`` is - # handled by the same code path. + # === Step 2. Destination-wise max including the physical null mass === + # The null initialization keeps empty and all-masked segments finite. group_max = xp_maximum_at( - xp.full((n_nodes, n_channel), float("-inf"), dtype=logits.dtype, device=device), + xp.zeros((n_nodes, n_channel), dtype=compute_dtype, device=device) + null_logit, dst, - logits_for_max, + effective_logits, ) # (N, n_channel) edge_max = xp.take(group_max, dst, axis=0) - zeros_en = xp.zeros((n_edge, n_channel), dtype=logits.dtype, device=device) - zeros_nn = xp.zeros((n_nodes, n_channel), dtype=logits.dtype, device=device) - edge_max = xp.where(xp.isfinite(edge_max), edge_max, zeros_en) - group_max_safe = xp.where(xp.isfinite(group_max), group_max, zeros_nn) - - # === Step 3. Envelope/SFPG-gated exponential terms === - exp_shifted = xp.exp(logits_2d - edge_max) - edge_weighted_exp = edge_weight_sq[:, None] * exp_shifted - - # === Step 4. Destination-wise normalization with positive denominator bias === - # Destination segment sum over ``dst`` (pt ``scatter_add``); invalid slots - # already carry zero weight. Layout-agnostic like the group max above. + + # === Step 3. Normalize edge and null masses in the shared shifted frame === + edge_exp = xp.exp(effective_logits - edge_max) + if source_ratio is not None: + # ``source_scale`` carries the forward magnitude in log space, while + # this linear ratio carries derivatives without differentiating + # ``log(src_weight)`` at extremely small positive gates. + edge_exp = edge_exp * source_ratio[:, None] denom_sum = xp_add_at( - xp.zeros((n_nodes, n_channel), dtype=logits.dtype, device=device), + xp.zeros((n_nodes, n_channel), dtype=compute_dtype, device=device), dst, - edge_weighted_exp, + edge_exp, ) # (N, n_channel) - denom = denom_sum + zeta * xp.exp(-group_max_safe) - - denom_edge = xp.take(denom, dst, axis=0) - alpha = edge_weighted_exp / (denom_edge + eps_f) + denominator = denom_sum + xp.exp(null_logit - group_max) + alpha = edge_exp / xp.take(denominator, dst, axis=0) + if promote: + alpha = xp.astype(alpha, input_dtype) return xp.reshape(alpha, (n_edge, n_focus, n_head)) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py index 6012cc4256..5ef1c92771 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py @@ -204,28 +204,20 @@ class C3CutoffEnvelope(NativeOP): Notes ----- - The envelope function is defined for scaled distance ``x = r / rcut`` as:: + For scaled distance ``x = r / rcut`` and ``u = 1 - x``, the envelope is + evaluated in the cancellation-free form:: - E(x) = 1 + x^p * (a + b*x + c*x^2 + d*x^3), for x < 1 - E(x) = 0, for x >= 1 + E_p(x) = u^4 * sum(comb(k + 3, 3) * x^k, k=0..p-1), for x < 1 + E_p(x) = 0, for x >= 1 - where the coefficients are chosen to satisfy:: + This positive-coefficient factorization satisfies:: E(0) = 1, E(1) = 0 E'(1) = 0, E''(1) = 0, E'''(1) = 0 - This ensures C^3 continuity at the cutoff boundary. The coefficients are:: + For the default exponent ``p=5``:: - a = -(p + 1)(p + 2)(p + 3) / 6 - b = p(p + 2)(p + 3) / 2 - c = -p(p + 1)(p + 3) / 2 - d = p(p + 1)(p + 2) / 6 - - For the default exponent p=5, the coefficients are a=-56, b=140, c=-120, - d=35:: - - E(x) = 1 + x^5 * (-56 + 140*x - 120*x^2 + 35*x^3) - = 1 - 56*x^5 + 140*x^6 - 120*x^7 + 35*x^8 + E_5(x) = u^4 * (1 + 4*x + 10*x^2 + 20*x^3 + 35*x^4) Parameters ---------- @@ -240,14 +232,6 @@ class C3CutoffEnvelope(NativeOP): Cutoff radius in Å. p : float Polynomial exponent. - a : float - Quadratic coefficient for x^p term. - b : float - Linear coefficient for x^(p+1) term. - c : float - Quadratic coefficient for x^(p+2) term. - d : float - Cubic coefficient for x^(p+3) term. """ def __init__( @@ -264,20 +248,25 @@ def __init__( self.rcut = float(rcut) self.p = int(exponent) self.precision = precision - self.coeff_a = -((self.p + 1) * (self.p + 2) * (self.p + 3)) / 6.0 - self.coeff_b = (self.p * (self.p + 2) * (self.p + 3)) / 2.0 - self.coeff_c = -(self.p * (self.p + 1) * (self.p + 3)) / 2.0 - self.coeff_d = (self.p * (self.p + 1) * (self.p + 2)) / 6.0 + self._series_coefficients = tuple( + float(math.comb(k + 3, 3)) for k in range(self.p) + ) def call(self, dst: Any) -> Any: """Compute the envelope value for given distances.""" xp = array_api_compat.array_namespace(dst) - d_scaled = xp.clip(dst / self.rcut, min=0.0, max=1.0) - poly = self.coeff_a + d_scaled * ( - self.coeff_b + d_scaled * (self.coeff_c + d_scaled * self.coeff_d) + device = array_api_compat.device(dst) + u = xp.clip((self.rcut - dst) / self.rcut, min=0.0, max=1.0) + x = 1.0 - u + series = xp.full( + x.shape, + self._series_coefficients[-1], + dtype=x.dtype, + device=device, ) - env_val = 1 + d_scaled**self.p * poly - return env_val * xp.astype(d_scaled < 1.0, dst.dtype) + for coefficient in reversed(self._series_coefficients[:-1]): + series = coefficient + x * series + return u**4 * series class InnerClamp(NativeOP): diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 178c43072e..53f945c446 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -853,6 +853,11 @@ def freeze_sezm_to_pt2( free by-product of the single backward, so exporting it carries no compute cost. """ + log.info( + "Set DP_TRITON_INFER to the desired level (0-3) before freezing; " + "the selected Triton inference kernels are baked into the .pt2 archive." + ) + from torch._inductor import ( aoti_compile_and_package, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/attention.py b/deepmd/pt/model/descriptor/sezm_nn/attention.py index 4f42188c2e..ec20c4ac93 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/attention.py +++ b/deepmd/pt/model/descriptor/sezm_nn/attention.py @@ -41,13 +41,13 @@ def segment_envelope_gated_softmax( Unconstrained denominator bias with shape (F, H). Softplus is applied to keep the bias strictly positive. eps - Small epsilon for denominator stability. + Small positive floor added to the physical null mass. src_weight Optional per-edge source-side multiplier with shape (E, 1) or - (E,). When provided the per-edge weight becomes - ``edge_env**2 * src_weight`` and the attention reduces to - ``edge_env**2 * src_weight * exp(logits) / - (zeta + sum(edge_env**2 * src_weight * exp(logits)))``. + (E,). When provided, the physical per-edge mass is + ``edge_env**2 * src_weight * exp(logits)`` and the denominator is the + sum of edge masses plus the positive null mass + ``softplus(z_bias_raw) + eps``. ``src_weight = 0`` therefore removes the source from both the numerator and the denominator, which is what SFPG needs so that a muted source does not even leak through the softmax @@ -60,65 +60,76 @@ def segment_envelope_gated_softmax( """ n_edge, n_focus, n_head = logits.shape n_channel = n_focus * n_head - eps_f = float(eps) + input_dtype = logits.dtype + compute_dtype = ( + torch.float32 if input_dtype in (torch.float16, torch.bfloat16) else input_dtype + ) - # === Step 1. Flatten (F, H) and build the effective per-edge weight === - logits_2d = logits.reshape(n_edge, n_channel) - edge_env_1d = edge_env.squeeze(-1).to(dtype=logits.dtype).clamp_min(0.0) - # edge_weight_sq acts as the non-negative multiplier applied to every - # ``exp(logit)`` term. Folding ``src_weight`` here guarantees that any - # edge with ``src_weight = 0`` is excluded from the group max, the - # numerator, and the denominator in a single pass. - edge_weight_sq = edge_env_1d.square() + # === Step 1. Build factor-wise effective logits === + # Computing the logarithms before multiplying the factors avoids losing a + # physically nonzero edge when ``edge_env**2 * src_weight`` underflows. + logits_2d = logits.reshape(n_edge, n_channel).to(dtype=compute_dtype) + edge_env_1d = edge_env.reshape(n_edge).to(dtype=compute_dtype) + edge_positive = edge_env_1d > 0.0 + ones = torch.ones_like(edge_env_1d) + log_weight = 2.0 * torch.log(torch.where(edge_positive, edge_env_1d, ones)) + active = edge_positive + source_ratio: torch.Tensor | None = None if src_weight is not None: - edge_weight_sq = edge_weight_sq * src_weight.reshape(n_edge).to( - dtype=logits.dtype - ).clamp_min(0.0) - zeta = F.softplus(z_bias_raw).reshape(1, n_channel).to(dtype=logits.dtype) - dst_index = dst.reshape(n_edge, 1).expand(n_edge, n_channel) - has_weight = edge_weight_sq > 0.0 - logits_for_max = torch.where( - has_weight.reshape(n_edge, 1), - logits_2d, + source_weight = src_weight.reshape(n_edge).to(dtype=compute_dtype) + source_positive = source_weight > 0.0 + safe_source = torch.where(source_positive, source_weight, ones) + source_scale = safe_source.detach() + log_weight = log_weight + torch.log(source_scale) + source_ratio = torch.where( + source_positive, + source_weight / source_scale, + torch.zeros_like(source_weight), + ) + active = active & source_positive + effective_logits = torch.where( + active.reshape(n_edge, 1), + logits_2d + log_weight.reshape(n_edge, 1), torch.full_like(logits_2d, float("-inf")), ) - # === Step 2. Destination-wise max for stable exponentials === - group_max = torch.full( - (n_nodes, n_channel), - float("-inf"), - dtype=logits.dtype, - device=logits.device, + null_mass = (F.softplus(z_bias_raw.to(dtype=compute_dtype)) + float(eps)).reshape( + 1, n_channel ) + null_logit = torch.log(null_mass) + dst_index = dst.reshape(n_edge, 1).expand(n_edge, n_channel) + + # === Step 2. Destination-wise max including the physical null mass === + # Initializing every segment with ``null_logit`` keeps empty and all-masked + # segments finite without a separate fallback branch. + group_max = null_logit.expand(n_nodes, n_channel).clone() group_max = torch.scatter_reduce( group_max, 0, dst_index, - logits_for_max, + effective_logits, reduce="amax", include_self=True, ) edge_max = group_max.index_select(0, dst) - edge_max = torch.where( - torch.isfinite(edge_max), edge_max, torch.zeros_like(edge_max) - ) - group_max_safe = torch.where( - torch.isfinite(group_max), group_max, torch.zeros_like(group_max) - ) - - # === Step 3. Envelope/SFPG-gated exponential terms === - exp_shifted = torch.exp(logits_2d - edge_max) - edge_weighted_exp = edge_weight_sq.reshape(n_edge, 1) * exp_shifted - # === Step 4. Destination-wise normalization with positive denominator bias === + # === Step 3. Normalize edge and null masses in the shared shifted frame === + edge_exp = torch.exp(effective_logits - edge_max) + if source_ratio is not None: + # ``source_scale`` carries the forward magnitude in log space, while + # this linear ratio carries derivatives without differentiating + # ``log(src_weight)`` at extremely small positive gates. + edge_exp = edge_exp * source_ratio.reshape(n_edge, 1) denom_sum = torch.zeros( n_nodes, n_channel, - dtype=logits.dtype, + dtype=compute_dtype, device=logits.device, ) - denom_sum = torch.scatter_add(denom_sum, 0, dst_index, edge_weighted_exp) - denom = denom_sum + zeta * torch.exp(-group_max_safe) + denom_sum = torch.scatter_add(denom_sum, 0, dst_index, edge_exp) + denominator = denom_sum + torch.exp(null_logit - group_max) + alpha = edge_exp / denominator.index_select(0, dst) - alpha = edge_weighted_exp / (denom.index_select(0, dst) + eps_f) + if alpha.dtype != input_dtype: + alpha = alpha.to(dtype=input_dtype) return alpha.reshape(n_edge, n_focus, n_head) diff --git a/deepmd/pt/model/descriptor/sezm_nn/radial.py b/deepmd/pt/model/descriptor/sezm_nn/radial.py index cef4e402c3..9622484e67 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/radial.py +++ b/deepmd/pt/model/descriptor/sezm_nn/radial.py @@ -199,28 +199,20 @@ class C3CutoffEnvelope(torch.nn.Module): Notes ----- - The envelope function is defined for scaled distance ``x = r / rcut`` as:: + For scaled distance ``x = r / rcut`` and ``u = 1 - x``, the envelope is + evaluated in the cancellation-free form:: - E(x) = 1 + x^p * (a + b*x + c*x^2 + d*x^3), for x < 1 - E(x) = 0, for x >= 1 + E_p(x) = u^4 * sum(comb(k + 3, 3) * x^k, k=0..p-1), for x < 1 + E_p(x) = 0, for x >= 1 - where the coefficients are chosen to satisfy:: + This positive-coefficient factorization satisfies:: E(0) = 1, E(1) = 0 E'(1) = 0, E''(1) = 0, E'''(1) = 0 - This ensures C^3 continuity at the cutoff boundary. The coefficients are:: + For the default exponent ``p=5``:: - a = -(p + 1)(p + 2)(p + 3) / 6 - b = p(p + 2)(p + 3) / 2 - c = -p(p + 1)(p + 3) / 2 - d = p(p + 1)(p + 2) / 6 - - For the default exponent p=5, the coefficients are a=-56, b=140, c=-120, - d=35:: - - E(x) = 1 + x^5 * (-56 + 140*x - 120*x^2 + 35*x^3) - = 1 - 56*x^5 + 140*x^6 - 120*x^7 + 35*x^8 + E_5(x) = u^4 * (1 + 4*x + 10*x^2 + 20*x^3 + 35*x^4) Parameters ---------- @@ -235,14 +227,6 @@ class C3CutoffEnvelope(torch.nn.Module): Cutoff radius in Å. p : float Polynomial exponent. - a : float - Quadratic coefficient for x^p term. - b : float - Linear coefficient for x^(p+1) term. - c : float - Quadratic coefficient for x^(p+2) term. - d : float - Cubic coefficient for x^(p+3) term. """ def __init__( @@ -261,44 +245,23 @@ def __init__( self.p = int(exponent) self.dtype = dtype self.device = env.DEVICE - coeff_a = -((self.p + 1) * (self.p + 2) * (self.p + 3)) / 6.0 - coeff_b = (self.p * (self.p + 2) * (self.p + 3)) / 2.0 - coeff_c = -(self.p * (self.p + 1) * (self.p + 3)) / 2.0 - coeff_d = (self.p * (self.p + 1) * (self.p + 2)) / 6.0 + self._series_coefficients = tuple( + float(math.comb(k + 3, 3)) for k in range(self.p) + ) self.register_buffer( "rcut_tensor", torch.tensor(self.rcut, dtype=self.dtype, device=self.device), persistent=False, ) - self.register_buffer( - "coeff_a", - torch.tensor(coeff_a, dtype=self.dtype, device=self.device), - persistent=False, - ) - self.register_buffer( - "coeff_b", - torch.tensor(coeff_b, dtype=self.dtype, device=self.device), - persistent=False, - ) - self.register_buffer( - "coeff_c", - torch.tensor(coeff_c, dtype=self.dtype, device=self.device), - persistent=False, - ) - self.register_buffer( - "coeff_d", - torch.tensor(coeff_d, dtype=self.dtype, device=self.device), - persistent=False, - ) def forward(self, dst: torch.Tensor) -> torch.Tensor: """Compute the envelope value for given distances.""" - d_scaled = (dst / self.rcut_tensor).clamp(min=0.0, max=1.0) - poly = self.coeff_a + d_scaled * ( - self.coeff_b + d_scaled * (self.coeff_c + d_scaled * self.coeff_d) - ) - env_val = 1 + d_scaled.pow(self.p) * poly - return env_val * ((d_scaled < 1.0).to(dst.dtype)) + u = ((self.rcut_tensor - dst) / self.rcut_tensor).clamp(min=0.0, max=1.0) + x = 1.0 - u + series = torch.full_like(x, self._series_coefficients[-1]) + for coefficient in reversed(self._series_coefficients[:-1]): + series = coefficient + x * series + return u.pow(4) * series class InnerClamp(nn.Module): diff --git a/examples/water/dpa4/README.md b/examples/water/dpa4/README.md index 3cbb0dde49..b832eeb071 100644 --- a/examples/water/dpa4/README.md +++ b/examples/water/dpa4/README.md @@ -7,7 +7,7 @@ water example dataset. The recommended model and descriptor type is `DPA4`; Input files: - `input.json`: baseline conservative energy training, using a compact - DPA4-Neo-style parameter set. + DPA4-Mini-style parameter set. - `input-zbl.json`: energy training with ZBL zone bridging. - `input_dens.json`: direct-force denoising training. - `input_multitask.json`: multitask training with a shared descriptor and diff --git a/examples/water/dpa4/input-zbl.json b/examples/water/dpa4/input-zbl.json index 14d675dcfc..ddd010233c 100644 --- a/examples/water/dpa4/input-zbl.json +++ b/examples/water/dpa4/input-zbl.json @@ -1,38 +1,37 @@ { - "_comment": "DPA4/SeZM energy-training example with ZBL zone bridging.", + "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", "model": { - "type": "DPA4", + "type": "dpa4", "type_map": [ "O", "H" ], "descriptor": { - "sel": 120, "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, - "lmax": 3, + "lmax": 2, "mmax": 1, "n_blocks": 2, - "so2_layers": 3, + "mixing_layers": 3, "radial_so2_mode": "degree_channel", "radial_so2_rank": 1, - "n_focus": 2, + "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": true, "ffn_neurons": 0, "ffn_so3_grid": true, "grid_mlp": false, - "grid_branch": 1, - "ffn_blocks": 2, - "sandwich_norm": [ - false, - true, - true, - false + "grid_branch": [ + 0, + 0, + 1 ], - "message_node_so3": true, + "ffn_blocks": 1, + "so3_readout": "mlp", "use_amp": true, "precision": "float32", "seed": 42 @@ -54,7 +53,7 @@ "type": "wsd", "start_lr": 4.5e-4, "stop_lr": 1e-6, - "warmup_steps": 5000, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2, "decay_phase_ratio": 0.65, "decay_type": "cosine" diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 1819a3afad..d9d49825d1 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -1,7 +1,7 @@ { - "_comment": "Baseline DPA4/SeZM energy-training example for the water dataset.", + "_comment": "DPA4-Mini energy-training example for the water dataset.", "model": { - "type": "DPA4", + "type": "dpa4", "type_map": [ "O", "H" @@ -12,38 +12,26 @@ "n_radial": 16, "edge_norm": false, "use_env_seed": true, - "edge_cartesian": false, - "node_cartesian": "none", - "lmax": 3, + "lmax": 2, "mmax": 1, "n_blocks": 2, "mixing_layers": 3, "radial_so2_mode": "degree_channel", "radial_so2_rank": 1, - "n_focus": 2, + "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": true, "ffn_neurons": 0, "ffn_so3_grid": true, - "so3_readout": "mlp", - "grid_mlp": [ - false, - false, - false - ], + "grid_mlp": false, "grid_branch": [ - 1, - 1, + 0, + 0, 1 ], - "ffn_blocks": 2, - "sandwich_norm": [ - false, - true, - true, - false - ], - "message_node_so3": true, + "ffn_blocks": 1, + "so3_readout": "mlp", "use_amp": true, "precision": "float32", "seed": 42 @@ -62,7 +50,7 @@ "type": "wsd", "start_lr": 4.5e-4, "stop_lr": 1e-6, - "warmup_steps": 5000, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2, "decay_phase_ratio": 0.65, "decay_type": "cosine" diff --git a/examples/water/dpa4/input_dens.json b/examples/water/dpa4/input_dens.json index 266f68eba7..b2ead2938b 100644 --- a/examples/water/dpa4/input_dens.json +++ b/examples/water/dpa4/input_dens.json @@ -1,38 +1,37 @@ { - "_comment": "DPA4/SeZM direct-force denoising training example.", + "_comment": "DPA4-Mini direct-force denoising training example.", "model": { - "type": "DPA4", + "type": "dpa4", "type_map": [ "O", "H" ], "descriptor": { - "sel": 120, "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, - "lmax": 3, + "lmax": 2, "mmax": 1, "n_blocks": 2, - "so2_layers": 3, + "mixing_layers": 3, "radial_so2_mode": "degree_channel", "radial_so2_rank": 1, - "n_focus": 2, + "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": true, "ffn_neurons": 0, "ffn_so3_grid": true, "grid_mlp": false, - "grid_branch": 1, - "ffn_blocks": 2, - "sandwich_norm": [ - false, - true, - true, - false + "grid_branch": [ + 0, + 0, + 1 ], - "message_node_so3": true, + "ffn_blocks": 1, + "so3_readout": "mlp", "use_amp": true, "precision": "float32", "seed": 42 @@ -51,7 +50,7 @@ "type": "wsd", "start_lr": 4.5e-4, "stop_lr": 1e-6, - "warmup_steps": 5000, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2, "decay_phase_ratio": 0.65, "decay_type": "cosine" diff --git a/examples/water/dpa4/input_multitask.json b/examples/water/dpa4/input_multitask.json index 15e97cdaba..90c38f74b3 100644 --- a/examples/water/dpa4/input_multitask.json +++ b/examples/water/dpa4/input_multitask.json @@ -1,5 +1,5 @@ { - "_comment": "DPA4/SeZM multitask example with a shared descriptor and case-conditioned shared fitting network.", + "_comment": "DPA4-Mini multitask example with a shared descriptor and case-conditioned shared fitting network.", "model": { "use_compile": false, "enable_tf32": true, @@ -9,32 +9,32 @@ "H" ], "descriptor": { - "sel": 120, + "type": "dpa4", "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, - "lmax": 3, + "lmax": 2, "mmax": 1, "n_blocks": 2, - "so2_layers": 3, + "mixing_layers": 3, "radial_so2_mode": "degree_channel", "radial_so2_rank": 1, - "n_focus": 2, + "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": true, "ffn_neurons": 0, "ffn_so3_grid": true, "grid_mlp": false, - "grid_branch": 1, - "ffn_blocks": 2, - "sandwich_norm": [ - false, - true, - true, - false + "grid_branch": [ + 0, + 0, + 1 ], - "message_node_so3": true, + "ffn_blocks": 1, + "so3_readout": "mlp", "use_amp": true, "precision": "float32", "seed": 42 @@ -52,7 +52,7 @@ }, "model_dict": { "water_1": { - "type": "DPA4", + "type": "dpa4", "type_map": "type_map", "descriptor": "descriptor", "fitting_net": "shared_fit_with_id", @@ -65,7 +65,7 @@ } }, "water_2": { - "type": "DPA4", + "type": "dpa4", "type_map": "type_map", "descriptor": "descriptor", "fitting_net": "shared_fit_with_id", @@ -82,7 +82,7 @@ "type": "wsd", "start_lr": 4.5e-4, "stop_lr": 1e-6, - "warmup_steps": 5000, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2, "decay_phase_ratio": 0.65, "decay_type": "cosine" diff --git a/examples/water/dpa4/lmp/README.md b/examples/water/dpa4/lmp/README.md index fdd31324d1..a62711fc87 100644 --- a/examples/water/dpa4/lmp/README.md +++ b/examples/water/dpa4/lmp/README.md @@ -6,12 +6,12 @@ same PyTorch implementation; DPA4 is the DPA-series user-facing name. ## Files -| File | Description | -| --------------- | -------------------------------------------------------------------------------------------------------------------------------- | -| `input.json` | Training configuration: tiny DPA4 / SeZM (`channels=16`, two blocks, fp32), 500 Adam steps on `examples/water/data/data_{0..3}`. | -| `pretrained.pt` | Shipped checkpoint for the LAMMPS smoke test. | -| `in.lammps` | 20-step NVT run at 330 K on 192 water molecules. | -| `water.lmp` | LAMMPS data file (192-atom liquid water cell). | +| File | Description | +| --------------- | ------------------------------------------------------------------------------------- | +| `input.json` | Demo DPA4/SeZM trained for 500 HybridMuon steps on `examples/water/data/data_{0..3}`. | +| `pretrained.pt` | Checkpoint produced from `input.json` for the LAMMPS smoke test. | +| `in.lammps` | 20-step NVT run at 330 K on 192 water molecules. | +| `water.lmp` | LAMMPS data file (192-atom liquid water cell). | The frozen `.pt2` archive is not included because AOTInductor packages are target-specific: they depend on the host's CPU/GPU, GPU compute @@ -41,16 +41,14 @@ Run the MD: lmp -in in.lammps ``` -Expected LAMMPS output: +The run should load the `.pt2` archive with a cutoff of 6 Å and two atom types, +then complete 20 steps with finite thermodynamic values. Exact values depend on +the trained checkpoint. ``` load model from: frozen_model.pt2 to gpu 0 rcut in model: 6 ntypes in model: 2 -Step PotEng KinEng TotEng Temp - 0 -29941.035 8.147 -29932.89 330.00 - 10 -29940.605 7.771 -29932.83 314.76 - 20 -29940.399 7.564 -29932.83 306.39 ``` ## Notes diff --git a/examples/water/dpa4/lmp/input.json b/examples/water/dpa4/lmp/input.json index 5a0d341a1a..0e0f42c755 100644 --- a/examples/water/dpa4/lmp/input.json +++ b/examples/water/dpa4/lmp/input.json @@ -7,51 +7,27 @@ "H" ], "descriptor": { - "type": "dpa4", - "sel": 120, "rcut": 6.0, - "env_exp": [ - 7, - 5 - ], "channels": 16, "n_radial": 6, - "radial_mlp": [ - 0 - ], + "edge_norm": false, "use_env_seed": true, - "random_gamma": true, - "l_schedule": [ - 1, - 1 - ], + "lmax": 1, "mmax": 1, - "so2_norm": false, - "so2_layers": 2, - "so2_attn_res": "none", + "n_blocks": 1, + "mixing_layers": 2, + "radial_so2_mode": "none", "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": false, "ffn_neurons": 0, + "ffn_so3_grid": false, "grid_mlp": false, + "grid_branch": 0, "ffn_blocks": 1, - "sandwich_norm": [ - true, - false, - true, - false - ], - "mlp_bias": false, - "layer_scale": false, - "full_attn_res": "none", - "block_attn_res": "none", - "s2_activation": [ - false, - true - ], - "activation_function": "silu", - "glu_activation": true, - "use_amp": false, + "so3_readout": "none", + "use_amp": true, "precision": "float32", "seed": 42 }, @@ -59,18 +35,17 @@ "neuron": [ 0 ], - "activation_function": "silu", "precision": "float32", "seed": 42 }, "use_compile": false, - "enable_tf32": false + "enable_tf32": true }, "learning_rate": { "type": "wsd", "start_lr": 0.0005, "stop_lr": 1e-06, - "warmup_steps": 50, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2, "decay_phase_ratio": 0.65, "decay_type": "cosine" diff --git a/examples/water/dpa4/lmp/pretrained.pt b/examples/water/dpa4/lmp/pretrained.pt index 6b15c4f6c0..de8bc4e055 100644 Binary files a/examples/water/dpa4/lmp/pretrained.pt and b/examples/water/dpa4/lmp/pretrained.pt differ diff --git a/examples/water/dpa4/lora_ft.json b/examples/water/dpa4/lora_ft.json index 7214133a77..b601fda2b2 100644 --- a/examples/water/dpa4/lora_ft.json +++ b/examples/water/dpa4/lora_ft.json @@ -1,38 +1,37 @@ { - "_comment": "DPA4/SeZM LoRA fine-tuning example.", + "_comment": "DPA4-Mini LoRA fine-tuning example.", "model": { - "type": "DPA4", + "type": "dpa4", "type_map": [ "O", "H" ], "descriptor": { - "sel": 120, "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, - "lmax": 3, + "lmax": 2, "mmax": 1, "n_blocks": 2, - "so2_layers": 3, + "mixing_layers": 3, "radial_so2_mode": "degree_channel", "radial_so2_rank": 1, - "n_focus": 2, + "n_focus": 1, "focus_dim": 0, "n_atten_head": 1, + "message_node_so3": true, "ffn_neurons": 0, "ffn_so3_grid": true, "grid_mlp": false, - "grid_branch": 1, - "ffn_blocks": 2, - "sandwich_norm": [ - false, - true, - true, - false + "grid_branch": [ + 0, + 0, + 1 ], - "message_node_so3": true, + "ffn_blocks": 1, + "so3_readout": "mlp", "use_amp": true, "precision": "float32", "seed": 42 @@ -55,7 +54,7 @@ "type": "cosine", "start_lr": 0.0005, "stop_lr": 1e-6, - "warmup_steps": 5000, + "warmup_ratio": 0.003, "warmup_start_factor": 0.2 }, "loss": { diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index 835d7c6a5d..a3b70be2f8 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -19,6 +19,7 @@ DescrptSeZM, ) from deepmd.pt.model.descriptor.sezm_nn import ( + C3CutoffEnvelope, DynamicRadialDegreeMixer, EdgeCartesianTensorProduct, ForceEmbedding, @@ -38,6 +39,7 @@ quaternion_multiply, quaternion_to_rotation_matrix, safe_norm, + segment_envelope_gated_softmax, ) from deepmd.pt.model.model import ( get_sezm_model, @@ -1828,6 +1830,137 @@ def test_invalid_params(self) -> None: InnerClamp(1.0, 1.0) +class TestCutoffNumerics(_SeZMTestCase): + """Numerical stability at cutoff-vanishing attention boundaries.""" + + def test_high_logit_edge_vanishes_continuously(self) -> None: + """A zero-weight high-logit edge must leave the segment continuously.""" + logits = torch.tensor( + [[[0.0]], [[20.0]]], dtype=torch.float64, device=self.device + ) + dst = torch.zeros(2, dtype=torch.int64, device=self.device) + z_bias_raw = torch.tensor( + [[math.log(math.expm1(1.0))]], dtype=torch.float64, device=self.device + ) + eps = 1.0e-7 + + def evaluate(crossing_envelope: float) -> torch.Tensor: + edge_env = torch.tensor( + [[1.0], [crossing_envelope]], + dtype=torch.float64, + device=self.device, + ) + return segment_envelope_gated_softmax( + logits, edge_env, dst, 1, z_bias_raw, eps + ) + + near = evaluate(1.0e-12) + zero = evaluate(0.0) + near_mass = math.exp(20.0) * 1.0e-24 + near_denominator = 2.0 + eps + near_mass + torch.testing.assert_close( + near[1, 0, 0], + torch.tensor( + near_mass / near_denominator, + dtype=torch.float64, + device=self.device, + ), + rtol=1.0e-12, + atol=0.0, + ) + expected_stable = 1.0 / (2.0 + eps) + torch.testing.assert_close( + near[0, 0, 0], + torch.tensor(expected_stable, dtype=torch.float64, device=self.device), + ) + torch.testing.assert_close( + zero[0, 0, 0], + torch.tensor(expected_stable, dtype=torch.float64, device=self.device), + ) + self.assertEqual(float(zero[1, 0, 0]), 0.0) + + def test_envelope_nextafter_cutoff_attention(self) -> None: + """Adjacent float32 distances must not create a spurious attention edge.""" + envelope_fn = C3CutoffEnvelope(6.0, exponent=5, dtype=torch.float32) + rcut = torch.tensor(6.0, dtype=torch.float32, device=self.device) + zero = torch.tensor(0.0, dtype=torch.float32, device=self.device) + r_near = torch.nextafter(rcut, zero) + r_inner = torch.nextafter(r_near, zero) + distances = torch.stack([r_near, r_inner])[:, None] + envelope = envelope_fn(distances) + + distance64 = distances[:, 0].to(torch.float64) + u = (6.0 - distance64) / 6.0 + x = 1.0 - u + reference = u**4 * (1.0 + x * (4.0 + x * (10.0 + x * (20.0 + 35.0 * x)))) + torch.testing.assert_close( + envelope[:, 0].to(torch.float64), + reference, + rtol=1.0e-6, + atol=0.0, + ) + self.assertTrue(bool((envelope >= 0.0).all())) + + logits = torch.tensor( + [[[0.0]], [[20.0]]], dtype=torch.float32, device=self.device + ) + dst = torch.zeros(2, dtype=torch.int64, device=self.device) + z_bias_raw = torch.tensor( + [[math.log(math.expm1(1.0))]], dtype=torch.float32, device=self.device + ) + for edge_envelope in envelope[:, 0]: + edge_env = torch.stack([torch.ones_like(edge_envelope), edge_envelope])[ + :, None + ] + alpha = segment_envelope_gated_softmax( + logits, edge_env, dst, 1, z_bias_raw, 1.0e-7 + ) + self.assertLess(float(alpha[1, 0, 0]), 1.0e-30) + + def test_tiny_source_weight_hessian(self) -> None: + """Log-domain source scaling must preserve the physical Hessian.""" + logits = torch.tensor( + [[[0.0]], [[20.0]]], dtype=torch.float32, device=self.device + ) + edge_env = torch.ones((2, 1), dtype=torch.float32, device=self.device) + dst = torch.zeros(2, dtype=torch.int64, device=self.device) + z_bias_raw = torch.tensor( + [[math.log(math.expm1(1.0))]], dtype=torch.float32, device=self.device + ) + eps = 1.0e-7 + + def attention_sum(source_weight: torch.Tensor) -> torch.Tensor: + return segment_envelope_gated_softmax( + logits, + edge_env, + dst, + 1, + z_bias_raw, + eps, + source_weight[:, None], + ).sum() + + null_mass = torch.nn.functional.softplus(z_bias_raw[0, 0]) + eps + + def physical_sum(source_weight: torch.Tensor) -> torch.Tensor: + edge_mass = source_weight * torch.exp(logits[:, 0, 0]) + return (edge_mass / (null_mass + edge_mass.sum())).sum() + + source_weight = torch.tensor( + [1.0, 1.0e-30], dtype=torch.float32, device=self.device + ) + actual = torch.autograd.functional.hessian(attention_sum, source_weight) + reference = torch.autograd.functional.hessian(physical_sum, source_weight) + self.assertTrue(bool(torch.isfinite(actual).all())) + torch.testing.assert_close( + actual[0, 0], + reference[0, 0], + rtol=1.0e-5, + atol=1.0e-6, + ) + torch.testing.assert_close(actual, reference, rtol=1.0e-5, atol=32.0) + + class TestEdgeNorm(_SeZMTestCase): """The ``edge_norm`` switch and its effect on cutoff smoothness. diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index c286f0fca7..feeba7d56d 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -8,6 +8,7 @@ ``assert_parity``) are added by the tasks that port each module. """ +import math import subprocess import sys @@ -2250,6 +2251,178 @@ def test_segment_envelope_gated_softmax(self, masked, use_src_weight) -> None: np.testing.assert_array_equal(alpha_dp[~valid], 0.0) assert np.all(np.isfinite(alpha_dp)) + def test_high_logit_edge_vanishes_continuously(self) -> None: + """A vanishing high-logit edge must agree with the physical limit.""" + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax as pt_softmax, + ) + + logits = np.array([[[0.0]], [[20.0]]], dtype=np.float64) + dst = np.zeros(2, dtype=np.int64) + z_bias_raw = np.array([[math.log(math.expm1(1.0))]], dtype=np.float64) + eps = 1.0e-7 + + def evaluate(crossing_envelope: float) -> np.ndarray: + edge_env = np.array([[1.0], [crossing_envelope]], dtype=np.float64) + alpha_dp = dp_softmax(logits, edge_env, dst, 1, z_bias_raw, eps) + alpha_pt = pt_softmax( + to_pt(logits), + to_pt(edge_env), + to_pt(dst), + 1, + to_pt(z_bias_raw), + eps, + ) + assert_parity(alpha_dp, alpha_pt) + return np.asarray(alpha_dp) + + near = evaluate(1.0e-12) + zero = evaluate(0.0) + near_mass = math.exp(20.0) * 1.0e-24 + near_denominator = 2.0 + eps + near_mass + np.testing.assert_allclose( + near[1, 0, 0], + near_mass / near_denominator, + rtol=1.0e-12, + atol=0.0, + ) + expected_stable = 1.0 / (2.0 + eps) + np.testing.assert_allclose( + near[0, 0, 0], expected_stable, rtol=1.0e-12, atol=0.0 + ) + np.testing.assert_allclose( + zero[0, 0, 0], expected_stable, rtol=1.0e-12, atol=0.0 + ) + assert zero[1, 0, 0] == 0.0 + + def test_envelope_nextafter_cutoff_attention(self) -> None: + """Adjacent float32 distances must remain stable in both implementations.""" + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as DPEnvelope, + ) + from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax as pt_softmax, + ) + from deepmd.pt.model.descriptor.sezm_nn.radial import ( + C3CutoffEnvelope as PTEnvelope, + ) + + rcut = np.float32(6.0) + zero = np.float32(0.0) + r_near = np.nextafter(rcut, zero) + r_inner = np.nextafter(r_near, zero) + distances = np.array([r_near, r_inner], dtype=np.float32)[:, None] + envelope_dp = np.asarray( + DPEnvelope(6.0, exponent=5, precision="float32").call(distances) + ) + envelope_pt = PTEnvelope(6.0, exponent=5, dtype=torch.float32)(to_pt(distances)) + + distance64 = distances[:, 0].astype(np.float64) + u = (6.0 - distance64) / 6.0 + x = 1.0 - u + reference = u**4 * (1.0 + x * (4.0 + x * (10.0 + x * (20.0 + 35.0 * x)))) + np.testing.assert_allclose( + envelope_dp[:, 0].astype(np.float64), + reference, + rtol=1.0e-6, + atol=0.0, + ) + assert_parity(envelope_dp, envelope_pt, rtol=1.0e-6, atol=0.0) + assert np.all(envelope_dp >= 0.0) + + logits = np.array([[[0.0]], [[20.0]]], dtype=np.float32) + dst = np.zeros(2, dtype=np.int64) + z_bias_raw = np.array([[math.log(math.expm1(1.0))]], dtype=np.float32) + for edge_envelope in envelope_dp[:, 0]: + edge_env = np.array([[1.0], [edge_envelope]], dtype=np.float32) + alpha_dp = np.asarray( + dp_softmax(logits, edge_env, dst, 1, z_bias_raw, 1.0e-7) + ) + alpha_pt = pt_softmax( + to_pt(logits), + to_pt(edge_env), + to_pt(dst), + 1, + to_pt(z_bias_raw), + 1.0e-7, + ) + assert_parity( + alpha_dp, + alpha_pt, + rtol=1.0e-5, + atol=float(np.finfo(np.float32).tiny), + ) + assert 0.0 <= alpha_dp[1, 0, 0] < 1.0e-30 + + def test_tiny_source_weight_hessian(self) -> None: + """The dpmodel and pt paths must preserve the physical Hessian.""" + from deepmd.dpmodel.descriptor.dpa4_nn.attention import ( + segment_envelope_gated_softmax as dp_softmax, + ) + from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax as pt_softmax, + ) + + logits = torch.tensor( + [[[0.0]], [[20.0]]], dtype=torch.float32, device=PT_DEVICE + ) + edge_env = torch.ones((2, 1), dtype=torch.float32, device=PT_DEVICE) + dst = torch.zeros(2, dtype=torch.int64, device=PT_DEVICE) + z_bias_raw = torch.tensor( + [[math.log(math.expm1(1.0))]], dtype=torch.float32, device=PT_DEVICE + ) + eps = 1.0e-7 + + def dp_attention_sum(source_weight: torch.Tensor) -> torch.Tensor: + return dp_softmax( + logits, + edge_env, + dst, + 1, + z_bias_raw, + eps, + source_weight[:, None], + ).sum() + + def pt_attention_sum(source_weight: torch.Tensor) -> torch.Tensor: + return pt_softmax( + logits, + edge_env, + dst, + 1, + z_bias_raw, + eps, + source_weight[:, None], + ).sum() + + null_mass = torch.nn.functional.softplus(z_bias_raw[0, 0]) + eps + + def physical_sum(source_weight: torch.Tensor) -> torch.Tensor: + edge_mass = source_weight * torch.exp(logits[:, 0, 0]) + return (edge_mass / (null_mass + edge_mass.sum())).sum() + + source_weight = torch.tensor( + [1.0, 1.0e-30], dtype=torch.float32, device=PT_DEVICE + ) + hessian_dp = torch.autograd.functional.hessian(dp_attention_sum, source_weight) + hessian_pt = torch.autograd.functional.hessian(pt_attention_sum, source_weight) + reference = torch.autograd.functional.hessian(physical_sum, source_weight) + assert bool(torch.isfinite(hessian_dp).all()) + torch.testing.assert_close( + hessian_dp[0, 0], + reference[0, 0], + rtol=1.0e-5, + atol=1.0e-6, + ) + torch.testing.assert_close(hessian_dp, hessian_pt, rtol=1.0e-5, atol=32.0) + torch.testing.assert_close(hessian_dp, reference, rtol=1.0e-5, atol=32.0) + def test_segment_softmax_arbitrary_degree(self) -> None: # The destination scatter is layout-agnostic: E need not be a multiple # of n_nodes and dst may carry an arbitrary (non-row-major) order with a