diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 81d47e45ea..896a348cfe 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -1007,10 +1007,14 @@ def _freeze_sezm_to_pt2( # Device-dependent Python branches resolve on the CPU tracing inputs, so # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2) - # contraction, while CUDA targets also bake the fused GIE scatter. + # contraction, while CUDA targets also bake the fused GIE scatter and + # permit level-2 SO(2) scalar capture despite the CPU tracing inputs. force_block_diag = target_device.type != "cpu" force_fused_scatter = target_device.type == "cuda" for module in model.modules(): + if isinstance(module, SO2Convolution): + module._scalar_capture_retains_rotation = True + module._force_scalar_capture_cuda_target = target_device.type == "cuda" if isinstance(module, SO2Linear): module._force_block_diag_matmul = force_block_diag if isinstance(module, GeometricInitialEmbedding): diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py index a02f428825..5377342d7d 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -188,6 +188,12 @@ def _rotate_mix_reference( dim = (lmax + 1) ** 2 n_deg = lmax + 1 reduced = 3 * lmax + 1 + if _has_no_edges(n_edge): + # opt_einsum 3.3 cannot plan zero-cost contractions (0 / 0). + # Empty reductions retain the zero derivative to every operand, + # including the nonempty node features and shared channel basis. + dependency = x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum() + return x.new_empty((n_focus, n_edge, reduced * focus_dim)) + dependency coeff = build_m_major_index(lmax, 1, device=x.device) d_to_m = wigner[:, :dim, :dim].index_select(1, coeff) x_local = torch.bmm(d_to_m, x.index_select(0, src)) # (E, reduced, C_wide) @@ -248,6 +254,15 @@ def _rotate_mix_backward_reference( dim = (lmax + 1) ** 2 n_deg = lmax + 1 reduced = 3 * lmax + 1 + if _has_no_edges(n_edge): + # All three edge cotangents are empty. Preserve their differentiable + # dependency without sending zero-size einsums to the path planner. + dependency = grad_u.sum() + x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum() + return ( + x.new_empty((n_edge, dim, c_wide)) + dependency, + wigner.new_empty(wigner.shape) + dependency, + kc.new_empty(kc.shape) + dependency, + ) coeff = build_m_major_index(lmax, 1, device=x.device) d_to_m = wigner[:, :dim, :dim].index_select(1, coeff) x_src = x.index_select(0, src) @@ -508,6 +523,7 @@ def _rotate_mix_fwd_kernel( kc_ptr, # (E, KSZ * RANK) compact kernel, or (E, L+1, CW) when RANK == 0 cb_ptr, # (RANK, CW) channel basis (unread when RANK == 0) u_ptr, # (F, E, ROW) focus-major output + scalar_ptr, # optional independent (E, F, CF), pitch 2*CF n_edge, x_sn, x_sd, @@ -516,6 +532,7 @@ def _rotate_mix_fwd_kernel( CW: tl.constexpr, # true C_wide; BC = next_power_of_2(CW) lanes with mask BC: tl.constexpr, RANK: tl.constexpr, + SAVE_SCALAR: tl.constexpr = False, ): """One program per edge, channels vectorized. @@ -576,6 +593,13 @@ def _rotate_mix_fwd_kernel( tl.float32 ) tl.store(u_ptr + f_off + o * CF, xl[o] * rad, mask=cmask) + if SAVE_SCALAR and o == 0: + scalar_off = ( + (chan // CF).to(tl.int64) * n_edge * (2 * CF) + + edge * (2 * CF) + + chan % CF + ) + tl.store(scalar_ptr + scalar_off, xl[o] * rad, mask=cmask) for o in tl.static_range(L): rad = tl.load(rad_base + (o + 1) * CW + chan, mask=cmask, other=0.0).to( tl.float32 @@ -606,6 +630,13 @@ def _rotate_mix_fwd_kernel( if RANK == 1: acc = acc * cb[0] tl.store(u_ptr + f_off + o * CF, acc, mask=cmask) + if SAVE_SCALAR and o == 0: + scalar_off = ( + (chan // CF).to(tl.int64) * n_edge * (2 * CF) + + edge * (2 * CF) + + chan % CF + ) + tl.store(scalar_ptr + scalar_off, acc, mask=cmask) for o in tl.static_range(L): accn = tl.zeros((BC,), dtype=tl.float32) accq = tl.zeros((BC,), dtype=tl.float32) @@ -2498,6 +2529,7 @@ def _rotate_mix_impl( kc, cb, u, + u, # unread unless SAVE_SCALAR is enabled n_edge, x.stride(0), x.stride(1), @@ -2506,12 +2538,82 @@ def _rotate_mix_impl( CW=c_wide, BC=triton.next_power_of_2(c_wide), RANK=int(rank), + SAVE_SCALAR=False, num_warps=warps, num_stages=stages, ) return u +def _rotate_mix_scalar_impl( + x: Tensor, + src: Tensor, + src_order: Tensor, + src_rowptr: Tensor, + wigner: Tensor, + kc: Tensor, + cb: Tensor, + lmax: int, + n_focus: int, + rank: int, +) -> tuple[Tensor, Tensor]: + """Emit the l=0 competition input without retaining a view of ``u``. + + The extra store uses the existing rotation accumulator; no second gather + or rotation is performed. The padded focus-major scalar allocation keeps + the competition reduction layout used by the original compiled path. + Its padding is never read. Both outputs own independent storage. + """ + del src_order, src_rowptr + n_edge = src.shape[0] + c_wide = int(x.shape[2]) + focus_dim = c_wide // int(n_focus) + row = (3 * int(lmax) + 1) * focus_dim + if not _use_triton(x): + u = _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank) + # CPU export tracing must describe the same independent, padded + # scalar storage as the CUDA implementation. + scalar = torch.empty_strided( + (n_edge, n_focus, focus_dim), + (2 * focus_dim, n_edge * 2 * focus_dim, 1), + device=x.device, + dtype=x.dtype, + ) + scalar.copy_(u[:, :, :focus_dim].permute(1, 0, 2)) + return u, scalar + u = torch.empty(n_focus, n_edge, row, device=x.device, dtype=x.dtype) + pitch = 2 * focus_dim + scalar = torch.empty_strided( + (n_edge, n_focus, focus_dim), + (pitch, n_edge * pitch, 1), + device=x.device, + dtype=x.dtype, + ) + if not _has_no_edges(n_edge): + warps, stages = rotate_mix_fwd_config(c_wide, int(lmax)) + wrap_triton(_rotate_mix_fwd_kernel)[(n_edge,)]( + x, + src, + wigner, + kc, + cb, + u, + scalar, + n_edge, + x.stride(0), + x.stride(1), + L=int(lmax), + CF=focus_dim, + CW=c_wide, + BC=triton.next_power_of_2(c_wide), + RANK=int(rank), + SAVE_SCALAR=True, + num_warps=warps, + num_stages=stages, + ) + return u, scalar + + def _rotate_mix_bwd_impl( grad_u: Tensor, x: Tensor, @@ -3608,6 +3710,9 @@ def _mixing_stack_train_bwd_impl( _rotate_mix_op = torch.library.triton_op( "sezm_triton::so2_rotate_mix", mutates_args=() )(_rotate_mix_impl) +_rotate_mix_scalar_op = torch.library.triton_op( + "sezm_triton::so2_rotate_mix_scalar", mutates_args=() +)(_rotate_mix_scalar_impl) _rotate_mix_bwd_op = torch.library.triton_op( "sezm_triton::so2_rotate_mix_bwd", mutates_args=() )(_rotate_mix_bwd_impl) @@ -3639,6 +3744,24 @@ def _mixing_stack_train_bwd_impl( )(_gated_act_bwd_impl) +@_rotate_mix_scalar_op.register_fake +def _rotate_mix_scalar_fake( + x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank +): + """Describe both outputs without specializing symbolic edges in einsum.""" + focus_dim = x.shape[2] // n_focus + n_edge = src.shape[0] + return ( + x.new_empty((n_focus, n_edge, (3 * lmax + 1) * focus_dim)), + torch.empty_strided( + (n_edge, n_focus, focus_dim), + (2 * focus_dim, n_edge * 2 * focus_dim, 1), + device=x.device, + dtype=x.dtype, + ), + ) + + @_rotate_mix_op.register_fake def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank): focus_dim = x.shape[2] // n_focus @@ -3892,6 +4015,25 @@ def _rotate_mix_backward(ctx, grad_u): return grad_x, None, None, None, grad_wigner, grad_kc, grad_cb, None, None, None +def _rotate_mix_scalar_setup_context(ctx, inputs, output): + _rotate_mix_setup_context(ctx, inputs, output) + ctx.set_materialize_grads(False) + + +def _rotate_mix_scalar_backward(ctx, grad_u, grad_scalar): + # Exactly the slice/permutation adjoint of the old competition branch. + # Keep pad/add visible to Inductor so it can fuse the gradient merge. + if grad_scalar is not None: + focus_dim = grad_scalar.shape[-1] + extra = torch.nn.functional.pad( + grad_scalar.permute(1, 0, 2), (0, 3 * ctx.lmax * focus_dim) + ) + grad_u = extra if grad_u is None else grad_u + extra + if grad_u is None: + return (None,) * 10 + return _rotate_mix_backward(ctx, grad_u) + + def _rotate_mix_bwd_setup_context(ctx, inputs, output): grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank = inputs ctx.save_for_backward(grad_u, x, src, wigner, kc, cb) @@ -4010,6 +4152,9 @@ def _segment_sum_backward(ctx, grad_out): _rotate_mix_op.register_autograd( _rotate_mix_backward, setup_context=_rotate_mix_setup_context ) +_rotate_mix_scalar_op.register_autograd( + _rotate_mix_scalar_backward, setup_context=_rotate_mix_scalar_setup_context +) _rotate_mix_bwd_op.register_autograd( _rotate_mix_bwd_backward, setup_context=_rotate_mix_bwd_setup_context ) @@ -5176,6 +5321,27 @@ def make_triton_rotate_mix(conv: SO2Convolution) -> _TritonRotateMix | None: ) +def _save_competition_scalar(conv: SO2Convolution, x: Tensor) -> bool: + """Capture only where competition backward retains the rotation storage. + + The PT-expt exporter can already release that storage without capture; + it pins this internal policy off before tracing. Runtime and PT export + retain the independent scalar buffer for the validated level-2 layout. + """ + return ( + getattr(conv, "_scalar_capture_retains_rotation", True) + and not conv.training + and conv.triton_infer_level == 2 + and conv.lmax == 3 + and conv.n_focus == 2 + and conv.so2_focus_dim == 32 + and conv.focus_compete + and x.dtype == torch.float32 + # Export traces CPU inputs before moving the graph to its target. + and (x.is_cuda or getattr(conv, "_force_scalar_capture_cuda_target", False)) + ) + + class _TritonSO2ValuePath: """Per-convolution entry running the SO(2) value path through the fused ops. @@ -5331,7 +5497,8 @@ def __call__( src_rowptr = torch.cat([counts.new_zeros(1), torch.cumsum(counts, 0)]) else: src_order, src_rowptr = csr - u0 = _rotate_mix_op( + save_scalar = _save_competition_scalar(conv, x) + rotation = (_rotate_mix_scalar_op if save_scalar else _rotate_mix_op)( x.contiguous(), src, src_order, @@ -5343,13 +5510,21 @@ def __call__( conv.n_focus, rank, ) + if save_scalar: + u0, scalar = rotation + else: + u0 = rotation # === Step 3. Cross-focus competition weight from the l = 0 scalars === apply_alpha = bool(conv.focus_compete and conv.n_focus > 1) if apply_alpha: # The small (E, F, Cf) copy keeps the softmax backward from # retaining a view of the whole focus-major activation. - gate_src = u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous() + gate_src = ( + scalar.contiguous() + if save_scalar + else u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous() + ) alpha = conv._focus_alpha(gate_src).to(u0.dtype).contiguous() else: alpha = torch.ones( diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index 44c0a91aac..6696b2e9ce 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -894,6 +894,7 @@ def launch(warps: int, stages: int) -> None: kc, cb, u, + u, # scalar output is disabled during the original rotate sweep n_edge, x.stride(0), x.stride(1), @@ -902,6 +903,7 @@ def launch(warps: int, stages: int) -> None: CW=c_wide, BC=triton.next_power_of_2(c_wide), RANK=1, + SAVE_SCALAR=False, num_warps=warps, num_stages=stages, ) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 8104cce43b..6572129c73 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1933,17 +1933,24 @@ def _trace_and_export_impl( # Device-dependent Python branches resolve on the CPU tracing inputs, so # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2) - # contraction, while CUDA targets also bake the fused GIE scatter. + # contraction, while CUDA targets also bake the fused GIE scatter and + # identify the CUDA target despite the CPU tracing inputs. This exporter + # already releases the full rotation output: a separate scalar store adds + # latency without reducing its peak, so keep the non-capture route. from deepmd.pt_expt.descriptor.dpa4_nn.embedding import ( GeometricInitialEmbedding, ) from deepmd.pt_expt.descriptor.dpa4_nn.so2 import ( + SO2Convolution, SO2Linear, ) force_block_diag = target_device.type != "cpu" force_fused_scatter = target_device.type == "cuda" for module in model.modules(): + if isinstance(module, SO2Convolution): + module._scalar_capture_retains_rotation = False + module._force_scalar_capture_cuda_target = target_device.type == "cuda" if isinstance(module, SO2Linear): module._force_block_diag_matmul = force_block_diag if isinstance(module, GeometricInitialEmbedding): diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py new file mode 100644 index 0000000000..54f822a36f --- /dev/null +++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Manual full-API CSC5 benchmark; run one route in each fresh process. + +Example: python source/tests/pt_expt/benchmark_so2_scalar_capture.py \ + --model model.pt --route CSC5 --cell 8 8 8 --out results/csc5 +""" +# ruff: noqa: T201 -- standalone benchmark CLI + +import argparse +import copy +import hashlib +import json +import os +import site +import subprocess +import sys +import time +import traceback +from pathlib import ( + Path, +) + +p = argparse.ArgumentParser(description=__doc__) +p.add_argument("--model", type=Path, required=True) +p.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[3]) +p.add_argument("--route", choices=("TORCH", "BASE", "CSC5"), required=True) +p.add_argument("--cell", type=int, nargs=3, default=(5, 5, 5)) +p.add_argument("--out", type=Path, required=True) +p.add_argument("--cache", type=Path) +p.add_argument("--gpu", type=int, default=0) +p.add_argument("--seed", type=int, default=0) +p.add_argument("--warmup", type=int, default=12) +p.add_argument("--samples", type=int, default=30) +p.add_argument("--legacy-pro", action="store_true") +a = p.parse_args() +assert a.warmup >= 3 and a.samples >= 10 +a.out.mkdir(parents=True, exist_ok=False) +cache = a.cache or a.out / "cache" + + +def write(name, data): + (a.out / name).write_text(json.dumps(data, indent=2, allow_nan=False)) + + +def sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +busy = subprocess.check_output( + [ + "nvidia-smi", + "-i", + str(a.gpu), + "--query-compute-apps=pid", + "--format=csv,noheader", + ], + text=True, +).strip() +if busy: + write("result.json", {"status": "busy", "pids": busy}) + raise SystemExit(3) +for key in list(os.environ): + if key.startswith("DP_"): + del os.environ[key] +os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) +os.environ.update( + CUDA_VISIBLE_DEVICES=str(a.gpu), + DP_COMPILE_INFER="1", + DP_COMPILE="0", + DP_TRITON_INFER="0" if a.route == "TORCH" else "2", + DP_CUDA_INFER="0", + DP_TF32_INFER="0", + DP_AMP_INFER="0", + OMP_NUM_THREADS="1", + DP_INTRA_OP_PARALLELISM_THREADS="1", + DP_INTER_OP_PARALLELISM_THREADS="1", + PYTORCH_ALLOC_CONF="expandable_segments:True", + TORCHINDUCTOR_COMPILE_THREADS="1", + TORCHINDUCTOR_CACHE_DIR=str(cache / "inductor"), + TRITON_CACHE_DIR=str(cache / "triton"), + TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD="1", +) +sys.path.insert(0, str(a.root.resolve())) +sys.meta_path[:] = [ + f + for f in sys.meta_path + if not type(f).__module__.startswith("_editable_skbc_deepmd_kit") +] +import deepmd + +assert Path(deepmd.__file__).resolve().parent.parent == a.root.resolve() +# Permit a source checkout to use its matching installed native extension. +for installed in site.getsitepackages(): + lib = Path(installed) / "deepmd" + if (lib / "lib").is_dir() and str(lib) not in deepmd.__path__: + deepmd.__path__.append(str(lib)) +import numpy as np +import torch +import triton +from ase.build import ( + bulk, +) +from compare_so2_scalar_capture import ( + configure_capture_route, +) + +from deepmd.calculator import ( + DP, +) +from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as scalar_path +from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError + +# BASE is an explicit benchmark-only ablation, not a public inference switch. +# On a pre-feature checkout BASE naturally uses its original implementation. +try: + configure_capture_route(a.route, scalar_path) +except ValueError as error: + write( + "result.json", {"status": "unsupported", "route": a.route, "reason": str(error)} + ) + raise SystemExit(2) from error + +torch.set_num_threads(1) +torch.set_float32_matmul_precision("highest") +torch.backends.cuda.matmul.allow_tf32 = False +torch.backends.cudnn.allow_tf32 = False +torch.use_deterministic_algorithms(False) +torch.cuda.reset_peak_memory_stats() +write( + "manifest.json", + { + "args": {k: str(v) for k, v in vars(a).items()}, + "model_sha256": sha(a.model), + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "revision": subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=a.root, text=True + ).strip(), + "source_sha256": sha( + a.root / "deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py" + ), + "benchmark_sha256": sha(__file__), + "flags": { + k: v + for k, v in os.environ.items() + if k.startswith(("DP_", "TORCHINDUCTOR", "TRITON_CACHE", "PYTORCH_ALLOC")) + }, + "gpu": subprocess.check_output( + [ + "nvidia-smi", + "-i", + str(a.gpu), + "--query-gpu=name,driver_version,uuid,clocks.sm,clocks.mem,temperature.gpu", + "--format=csv", + ], + text=True, + ), + }, +) +result = { + "status": "running", + "route": a.route, + "n_atoms": int(8 * np.prod(a.cell)), + "atol": 1e-5, + "rtol": 1e-4, +} +arrays = {} +started = time.time() +try: + atoms = bulk("C", "diamond", a=3.567, cubic=True).repeat(tuple(a.cell)) + rng = np.random.default_rng(a.seed) + frames = [] + for _ in range(3): + frames.append( + atoms.positions + 0.02 * rng.standard_normal(atoms.positions.shape) + ) + frames = np.stack(frames) + np.savez_compressed( + a.out / "inputs.npz", + positions=frames, + cell=np.asarray(atoms.cell), + numbers=atoms.numbers, + ) + result["input_sha256"] = sha(a.out / "inputs.npz") + load = torch.load + if a.legacy_pro: + assert ( + sha(a.model) + == "446a6e893101fc52190cfccabb9d2328863d651aca36981fd8d3401c71a61221" + ) + + def legacy_load(file, *args, **kwargs): + obj = load(file, *args, **kwargs) + if ( + isinstance(file, (str, Path)) + and Path(file).resolve() == a.model.resolve() + ): + state = obj["model"] if "model" in obj else obj + state["_extra_state"] = copy.deepcopy(state["_extra_state"]) + fitting = state["_extra_state"]["model_params"]["fitting_net"] + assert fitting["type"] == "ener" + fitting["type"] = "dpa4_ener" + return obj + + torch.load = legacy_load + result["legacy_pro"] = ( + "in-memory fitting type alias ener -> dpa4_ener; weights unchanged" + ) + try: + calc = DP(model=str(a.model), freeze=False) + finally: + torch.load = load + assert not calc.dp.get_model().training + result["parameters"] = sum(t.numel() for t in calc.dp.get_model().parameters()) + write("model_definition.json", calc.dp.get_model_def_script()) + + def evaluate(pos): + calc.reset() + sample = atoms.copy() + sample.positions = pos + sample.calc = calc + forces = sample.get_forces().copy() + return { + "energy": np.asarray(sample.get_potential_energy()), + "forces": forces, + "virial": np.asarray(calc.results["virial"]).copy(), + "stress": sample.get_stress().copy(), + } + + def peak(): + return { + "allocated": torch.cuda.max_memory_allocated(), + "reserved": torch.cuda.max_memory_reserved(), + } + + def stats(values): + return { + "median": float(np.median(values)), + "p90": float(np.percentile(values, 90)), + "min": min(values), + "max": max(values), + "samples": values, + } + + result["accuracy_wall_ms"] = [] + for repeat in range(3): + for frame, pos in enumerate(frames): + tick = time.perf_counter() + out = evaluate(pos) + torch.cuda.synchronize() + assert all(np.isfinite(v).all() for v in out.values()) + result["accuracy_wall_ms"].append(1000 * (time.perf_counter() - tick)) + arrays.update({f"{repeat}_{frame}_{k}": v for k, v in out.items()}) + np.savez_compressed(a.out / "outputs.npz", **arrays) + if repeat == frame == 0: + result["first_peak"] = peak() + print(a.route, len(atoms), repeat, frame, flush=True) + for i in range(a.warmup): + evaluate(frames[i % 3]) + wall, events = [], [] + for i in range(a.samples): + torch.cuda.synchronize() + begin, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + tick = time.perf_counter() + begin.record() + out = evaluate(frames[i % 3]) + end.record() + torch.cuda.synchronize() + wall.append(1000 * (time.perf_counter() - tick)) + events.append(begin.elapsed_time(end)) + assert all(np.isfinite(v).all() for v in out.values()) + full = peak() + torch.cuda.reset_peak_memory_stats() + for pos in frames: + evaluate(pos) + torch.cuda.synchronize() + result["steady_peak"] = peak() + result["complete_peak"] = {k: max(full[k], result["steady_peak"][k]) for k in full} + from torch._inductor.codecache import ( + PyCodeCache, + ) + + wrappers = [ + str(m.__file__) + for m in PyCodeCache.modules + if "def call(" in Path(m.__file__).read_text() + ] + assert wrappers + result.update( + status="completed", + finite=True, + wall_ms=stats(wall), + event_ms=stats(events), + wrappers=wrappers, + ) +except Exception as exc: + result.update( + status="oom" + if isinstance(exc, (torch.OutOfMemoryError, DPOutOfMemoryError)) + else "failed", + exception=traceback.format_exc(), + complete_peak={ + "allocated": torch.cuda.max_memory_allocated(), + "reserved": torch.cuda.max_memory_reserved(), + }, + ) + print(result["exception"], flush=True) +finally: + result["elapsed_s"] = time.time() - started + write("result.json", result) +raise SystemExit(0 if result["status"] == "completed" else 1) diff --git a/source/tests/pt_expt/compare_so2_scalar_capture.py b/source/tests/pt_expt/compare_so2_scalar_capture.py new file mode 100644 index 0000000000..9b0b6e24c1 --- /dev/null +++ b/source/tests/pt_expt/compare_so2_scalar_capture.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compare same-frame repeats from benchmark_so2_scalar_capture.py. + +Run with --torch DIR --base DIR --candidate DIR. An OOM reference is unavailable, +not a pass. Output is JSON; absolute units follow the ASE energy/force/virial/stress +arrays. Relative errors use max(abs(reference), 1e-8) as the denominator. +The worst tolerance ratio records its frame, repeat pair and component values. +""" +# ruff: noqa: T201 -- standalone comparison CLI + +import argparse +import json +from itertools import ( + product, +) +from pathlib import ( + Path, +) + +import numpy as np + + +def configure_capture_route(route, scalar_path): + """Reject missing CSC5 support; retain supported-layout fallback behavior.""" + available = hasattr(scalar_path, "_save_competition_scalar") + if route == "CSC5" and not available: + raise ValueError("CSC5 is unavailable in the selected source checkout") + if route == "BASE" and available: + scalar_path._save_competition_scalar = lambda conv, x: False + + +def compare(candidate, reference): + """Compare all 3x3 repeat pairs for each of the three corresponding frames.""" + fields = {} + for field in ("energy", "forces", "virial", "stress"): + rows = [] + for frame, i, j in product(range(3), range(3), range(3)): + x = candidate[f"{i}_{frame}_{field}"] + y = reference[f"{j}_{frame}_{field}"] + assert x.shape == y.shape + delta = np.abs(x - y) + ratio = delta / (1e-5 + 1e-4 * abs(y)) + index = np.unravel_index(np.argmax(ratio), ratio.shape) + rows.append( + { + "max_abs": float(delta.max()), + "mean_abs": float(delta.mean()), + "max_relative": float((delta / np.maximum(abs(y), 1e-8)).max()), + "max_ratio": float(ratio.max()), + "worst": { + "frame": frame, + "candidate_repeat": i, + "reference_repeat": j, + "index": [int(k) for k in index], + "candidate": float(x[index]), + "reference": float(y[index]), + "abs_error": float(delta[index]), + "tolerance": float(1e-5 + 1e-4 * abs(y[index])), + }, + "finite": bool(np.isfinite(x).all() and np.isfinite(y).all()), + } + ) + fields[field] = { + k: max(row[k] for row in rows) + for k in ("max_abs", "mean_abs", "max_relative", "max_ratio") + } + fields[field]["passed"] = all( + row["finite"] and row["max_ratio"] <= 1 for row in rows + ) + fields[field]["worst"] = max(rows, key=lambda row: row["max_ratio"])["worst"] + return fields + + +def validate_identity(result, reference, manifest, reference_manifest): + """Reject cross-checkpoint or cross-input comparisons before scoring.""" + for key, left, right in ( + ("input_sha256", result, reference), + ("model_sha256", manifest, reference_manifest), + ): + if not left.get(key) or left[key] != right.get(key): + raise ValueError(f"Comparison requires matching {key}") + + +def main(): + """Load benchmark routes and report compatible repeated comparisons.""" + parser = argparse.ArgumentParser(description=__doc__) + for route in ("torch", "base", "candidate"): + parser.add_argument(f"--{route}", type=Path, required=True) + args = parser.parse_args() + results, arrays, out, manifests = {}, {}, {}, {} + for route, path in vars(args).items(): + results[route] = json.loads((path / "result.json").read_text()) + out[route] = {"status": results[route]["status"]} + if results[route]["status"] == "completed": + manifests[route] = json.loads((path / "manifest.json").read_text()) + arrays[route] = dict(np.load(path / "outputs.npz")) + out[route]["repeat"] = compare(arrays[route], arrays[route]) + for route in arrays: + for reference in ("torch", "base"): + if route != reference and reference in arrays: + validate_identity( + results[route], + results[reference], + manifests[route], + manifests[reference], + ) + out[route][f"vs_{reference}"] = compare( + arrays[route], arrays[reference] + ) + print(json.dumps(out, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py new file mode 100644 index 0000000000..e549917039 --- /dev/null +++ b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Independent scalar storage emitted by the existing rotation producer.""" + +from types import ( + SimpleNamespace, +) + +import pytest +import torch + +from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not vp.SO2_VALUE_PATH_TRITON_AVAILABLE, + reason="CUDA and Triton are required", +) + + +@pytest.fixture(autouse=True) +def _strict_fp32(): + old = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + yield + torch.backends.cuda.matmul.allow_tf32 = old + + +def _inputs(edges: int, rank: int, zero: bool = False) -> tuple: + torch.manual_seed(42) + # Noncontiguous node stride; the existing kernel requires channel stride 1. + x = (torch.randn(14, 16, 64, device="cuda") * 0.1)[::2].detach() + if zero: + x.zero_() + x.requires_grad_() + src = torch.arange(edges, device="cuda") % 7 + order = src.argsort(stable=True) + counts = torch.bincount(src, minlength=7) + ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0))) + w = torch.randn(edges, 16, 16, device="cuda") * 0.1 + mask = torch.zeros(16, 16, device="cuda") + for degree in range(4): + mask[degree**2 : (degree + 1) ** 2, degree**2 : (degree + 1) ** 2] = 1 + w = (w * mask).requires_grad_() + shape = (edges, 4, 64) if rank == 0 else (edges, 25 * rank) + kc = (torch.randn(shape, device="cuda") * 0.1).requires_grad_() + cb = torch.randn(1 if rank == 0 else rank * 64, device="cuda") * 0.1 + cb.requires_grad_(rank > 0) + return x, src, order, ptr, w, kc, cb, 3, 2, rank + + +@pytest.mark.parametrize("edges,rank", [(0, 1), (1, 0), (17, 1), (257, 2)]) +@pytest.mark.parametrize("branch", ["both", "main", "scalar"]) +@pytest.mark.parametrize("zero", [False, True]) +def test_capture_and_vjp(edges: int, rank: int, branch: str, zero: bool) -> None: + args = _inputs(edges, rank, zero) + original = vp._rotate_mix_op(*args) + expected = (original, original[:, :, :32].permute(1, 0, 2)) + actual = vp._rotate_mix_scalar_op(*args) + for a, b in zip(actual, expected, strict=True): + torch.testing.assert_close(a, b, rtol=0, atol=0) + if edges: + assert ( + actual[0].untyped_storage().data_ptr() + != actual[1].untyped_storage().data_ptr() + ) + assert actual[1].stride() == (64, edges * 64, 1) + chosen = {"both": (0, 1), "main": (0,), "scalar": (1,)}[branch] + upstream = tuple(torch.randn_like(expected[i]) for i in chosen) + leaves = (args[0], args[4], args[5]) + ((args[6],) if rank else ()) + ref_grad = torch.autograd.grad(tuple(expected[i] for i in chosen), leaves, upstream) + got_grad = torch.autograd.grad(tuple(actual[i] for i in chosen), leaves, upstream) + for a, b in zip(got_grad, ref_grad, strict=True): + assert torch.isfinite(a).all() + torch.testing.assert_close(a, b, rtol=1e-4, atol=1e-5) + + +def test_vjp_can_be_differentiated() -> None: + """Preserve the existing rotation's differentiable backward contract.""" + args = _inputs(17, 1, False) + original = vp._rotate_mix_op(*args) + reference = (original, original[:, :, :32].permute(1, 0, 2)) + actual = vp._rotate_mix_scalar_op(*args) + + def derivatives(outputs): + loss = sum(v.square().sum() for v in outputs) + first = torch.autograd.grad(loss, args[0], create_graph=True)[0] + second = torch.autograd.grad(first.square().sum(), args[0])[0] + return first, second + + expected_grads = derivatives(reference) + actual_grads = derivatives(actual) + for actual_grad, expected_grad in zip(actual_grads, expected_grads, strict=True): + assert torch.isfinite(actual_grad).all() + torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-4, atol=1e-5) + + +def test_float64_reference() -> None: + args = _inputs(17, 2) + reference_args = tuple( + v.detach().double().requires_grad_(v.requires_grad) + if isinstance(v, torch.Tensor) and v.is_floating_point() + else v + for v in args + ) + # Wigner's off-degree entries are structural zeros, not independent + # variables. Project the dense reference onto the same tangent space. + block_mask = torch.zeros(16, 16, device="cuda", dtype=torch.float64) + for degree in range(4): + block_mask[degree**2 : (degree + 1) ** 2, degree**2 : (degree + 1) ** 2] = 1 + ref = vp._rotate_mix_reference( + reference_args[0], + reference_args[1], + reference_args[4] * block_mask, + reference_args[5], + reference_args[6], + 3, + 2, + 2, + ) + got, scalar = vp._rotate_mix_scalar_op(*args) + torch.testing.assert_close(got.double(), ref, rtol=1e-4, atol=1e-5) + g, gs = torch.randn_like(got), torch.randn_like(scalar) + actual = torch.autograd.grad( + (got, scalar), (args[0], args[4], args[5], args[6]), (g, gs) + ) + expected = torch.autograd.grad( + (ref, ref[:, :, :32].permute(1, 0, 2)), + tuple(reference_args[i] for i in (0, 4, 5, 6)), + (g.double(), gs.double()), + ) + for a, b in zip(actual, expected, strict=True): + torch.testing.assert_close(a.double(), b, rtol=1e-4, atol=1e-5) + + +def test_schema_and_fake() -> None: + torch.library.opcheck( + vp._rotate_mix_scalar_op, + _inputs(17, 1), + test_utils=("test_schema", "test_faketensor", "test_autograd_registration"), + ) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_dynamic_compile(capture: bool) -> None: + op = vp._rotate_mix_scalar_op if capture else vp._rotate_mix_op + compiled = torch.compile(op, fullgraph=True, dynamic=True) + for edges in (17, 33): + args = _inputs(edges, 1) + actual = compiled(*args) + expected = op(*args) + if not capture: + actual, expected = (actual,), (expected,) + for a, b in zip(actual, expected, strict=True): + torch.testing.assert_close(a, b, rtol=0, atol=0) + g = tuple(torch.randn_like(t) for t in expected) + leaves = tuple(args[i] for i in (0, 4, 5)) + a = torch.autograd.grad(actual, leaves, g) + b = torch.autograd.grad(expected, leaves, g) + for x, y in zip(a, b, strict=True): + torch.testing.assert_close(x, y, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize( + "field,value", + [ + ("training", True), + ("triton_infer_level", 1), + ("triton_infer_level", 3), + ("lmax", 2), + ("n_focus", 1), + ("so2_focus_dim", 64), + ("focus_compete", False), + ("_scalar_capture_retains_rotation", False), + ], +) +def test_scope(monkeypatch: pytest.MonkeyPatch, field: str, value: object) -> None: + conv = SimpleNamespace( + training=False, + triton_infer_level=2, + lmax=3, + n_focus=2, + so2_focus_dim=32, + focus_compete=True, + ) + x = torch.empty(1, 16, 64, device="cuda") + assert vp._save_competition_scalar(conv, x) + assert not vp._save_competition_scalar(conv, x.cpu()) + assert not vp._save_competition_scalar(conv, x.double()) + setattr(conv, field, value) + assert not vp._save_competition_scalar(conv, x) diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py new file mode 100644 index 0000000000..2abf9ffce9 --- /dev/null +++ b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Scalar capture must be selected and consumed, not merely numerically valid.""" + +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + patch, +) + +import pytest +import torch + +from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + +GPU = pytest.mark.skipif( + not torch.cuda.is_available() or not vp.SO2_VALUE_PATH_TRITON_AVAILABLE, + reason="CUDA and Triton are required", +) + + +@pytest.fixture +def strict_environment(monkeypatch): + """Keep the tested route independent of ambient acceleration flags.""" + for name, value in { + "DP_TRITON_INFER": "2", + "DP_CUDA_INFER": "0", + "DP_TRITON_TRAIN": "0", + "DP_CUDA_TRAIN": "0", + "DP_CUTE_INFER": "0", + "DP_CUTILE_INFER": "0", + }.items(): + monkeypatch.setenv(name, value) + old = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + yield + torch.backends.cuda.matmul.allow_tf32 = old + + +def test_export_target_scope(): + """Target and retained-storage policy both constrain CPU tracing.""" + conv = SimpleNamespace( + training=False, + triton_infer_level=2, + lmax=3, + n_focus=2, + so2_focus_dim=32, + focus_compete=True, + ) + x = torch.empty(7, 16, 64, device="cpu") + assert not vp._save_competition_scalar(conv, x) + conv._force_scalar_capture_cuda_target = True + assert vp._save_competition_scalar(conv, x) + conv._force_scalar_capture_cuda_target = False + assert not vp._save_competition_scalar(conv, x) + conv._force_scalar_capture_cuda_target = True + conv._scalar_capture_retains_rotation = False + assert not vp._save_competition_scalar(conv, x) + + +@pytest.mark.parametrize( + "field,value", + [ + ("training", True), + ("triton_infer_level", 0), + ("triton_infer_level", 1), + ("triton_infer_level", 3), + ("lmax", 2), + ("n_focus", 1), + ("so2_focus_dim", 64), + ("focus_compete", False), + ], +) +def test_cpu_tracing_preserves_selection_constraints(field, value): + """A CUDA export target must not bypass inference/layout eligibility.""" + conv = SimpleNamespace( + training=False, + triton_infer_level=2, + lmax=3, + n_focus=2, + so2_focus_dim=32, + focus_compete=True, + _force_scalar_capture_cuda_target=True, + ) + x = torch.empty(7, 16, 64, device="cpu") + assert vp._save_competition_scalar(conv, x) + assert not vp._save_competition_scalar(conv, x.double()) + setattr(conv, field, value) + assert not vp._save_competition_scalar(conv, x) + + +@pytest.mark.parametrize("edges", [0, 1, 17]) +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_cpu_scalar_layout_and_vjp(edges, rank): + """CPU reference and fake/CUDA layouts share the same output contract.""" + torch.manual_seed(42) + x = torch.randn(7, 16, 64, device="cpu", requires_grad=True) + src = torch.arange(edges, device="cpu") % 7 + order = src.argsort(stable=True) + ptr = torch.cat((src.new_zeros(1), torch.bincount(src, minlength=7).cumsum(0))) + w = torch.eye(16, device="cpu").repeat(edges, 1, 1).requires_grad_() + kc_shape = (edges, 4, 64) if rank == 0 else (edges, 25 * rank) + kc = torch.randn(kc_shape, device="cpu", requires_grad=True) + cb = torch.randn( + 1 if rank == 0 else 64 * rank, device="cpu", requires_grad=rank > 0 + ) + args = (x, src, order, ptr, w, kc, cb, 3, 2, rank) + u, scalar = vp._rotate_mix_scalar_op(*args) + ref = vp._rotate_mix_op(*args) + torch.testing.assert_close(u, ref, atol=0, rtol=0) + torch.testing.assert_close(scalar, ref[:, :, :32].permute(1, 0, 2), atol=0, rtol=0) + assert scalar.stride() == (64, edges * 64, 1) + if edges: + assert scalar.untyped_storage().data_ptr() != u.untyped_storage().data_ptr() + gu, gs = torch.randn_like(u), torch.randn_like(scalar) + leaves = (x, w, kc) + ((cb,) if rank else ()) + actual = torch.autograd.grad((u, scalar), leaves, (gu, gs)) + expected = torch.autograd.grad( + (ref, ref[:, :, :32].permute(1, 0, 2)), leaves, (gu, gs) + ) + for a, b in zip(actual, expected, strict=True): + torch.testing.assert_close(a, b, atol=1e-5, rtol=1e-4) + assert torch.isfinite(a).all() + if not edges: + torch.testing.assert_close(a, torch.zeros_like(a), atol=0, rtol=0) + + +@pytest.mark.parametrize("branch", ["both", "main", "scalar"]) +def test_empty_cpu_second_derivative(branch): + """Empty output branches retain zero first and second derivatives.""" + x = torch.randn(3, 16, 64, device="cpu", requires_grad=True) + src = torch.empty(0, device="cpu", dtype=torch.long) + ptr = torch.zeros(4, device="cpu", dtype=torch.long) + w = torch.empty(0, 16, 16, device="cpu", requires_grad=True) + kc = torch.empty(0, 25, device="cpu", requires_grad=True) + cb = torch.randn(64, device="cpu", requires_grad=True) + outputs = vp._rotate_mix_scalar_op(x, src, src, ptr, w, kc, cb, 3, 2, 1) + chosen = {"both": outputs, "main": outputs[:1], "scalar": outputs[1:]}[branch] + leaves = (x, w, kc, cb) + first = torch.autograd.grad( + sum(v.square().sum() for v in chosen), leaves, create_graph=True + ) + second = torch.autograd.grad(sum(v.square().sum() for v in first), leaves) + for gradient in (*first, *second): + torch.testing.assert_close(gradient, torch.zeros_like(gradient), rtol=0, atol=0) + + +@GPU +def test_capture_freeze_load(monkeypatch, strict_environment, tmp_path): + """Compile a real capture-enabled package and compare loaded predictions.""" + import re + import zipfile + + import deepmd.pt_expt.utils.env as export_env + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt.entrypoints import ( + freeze_pt2, + ) + from deepmd.pt.model.model import ( + get_model, + ) + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + + # The launch-table sweep is a performance concern; retain the real + # tracing, torch.export, AOTI compilation and package loader below. + monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None) + monkeypatch.setattr(export_env, "DEVICE", torch.device("cuda:0")) + params = _model_params() + params["descriptor"].update(radial_so2_mode="degree_channel", radial_so2_rank=1) + model = get_model(params).to("cuda").eval() + model.requires_grad_(False) + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.2, 0.2, 0.1], [0.3, 1.1, 0.2], [0.1, 0.2, 1.3]]], + device="cuda", + dtype=torch.float64, + ) + atype = torch.zeros(1, 4, device="cuda", dtype=torch.long) + reference = model(coord, atype, None) + checkpoint, package = tmp_path / "model.pt", tmp_path / "model.pt2" + wrapper = ModelWrapper(model, model_params=params) + torch.save({"model": wrapper.state_dict()}, checkpoint) + freeze_pt2.freeze_sezm_to_pt2( + str(checkpoint), + str(package), + device=torch.device("cuda:0"), + with_comm=False, + ) + with zipfile.ZipFile(package) as archive: + wrappers = [ + archive.read(name).decode() + for name in archive.namelist() + if name.endswith(".wrapper.cpp") + ] + assert any( + re.search(r"['\"]SAVE_SCALAR['\"]\s*:\s*True", text) for text in wrappers + ) + evaluator = DeepPot(str(package)) + energy, force, virial = evaluator.eval( + coord.detach().cpu().numpy(), + None, + [0] * 4, + ) + for actual, name in ((energy, "energy"), (force, "force"), (virial, "virial")): + expected = reference[name].detach().cpu() + torch.testing.assert_close( + torch.as_tensor(actual).reshape(expected.shape), + expected, + rtol=1e-4, + atol=1e-5, + ) + + +@GPU +@pytest.mark.parametrize("backend", ["pt", "pt_expt"]) +@pytest.mark.parametrize("retains_rotation", [False, True]) +def test_value_path_dispatch( + monkeypatch, strict_environment, backend, retains_rotation +): + """Exercise the production callable and prove the competition data source.""" + if backend == "pt": + from deepmd.pt.model.descriptor.sezm_nn.so2 import ( + SO2Convolution, + ) + + precision = {"dtype": torch.float32} + else: + from deepmd.pt_expt.descriptor.dpa4_nn.so2 import ( + SO2Convolution, + ) + + precision = {"precision": "float32"} + conv = ( + SO2Convolution( + lmax=3, + mmax=1, + channels=32, + n_focus=2, + focus_dim=32, + mixing_layers=3, + radial_so2_mode="degree_channel", + radial_so2_rank=1, + n_atten_head=1, + seed=7, + trainable=False, + **precision, + ) + .to("cuda") + .eval() + ) + path = vp.make_triton_value_path(conv) + assert path is not None + conv._scalar_capture_retains_rotation = retains_rotation + capture = retains_rotation + x = torch.randn(7, 16, 64, device="cuda") * 0.1 + radial = torch.randn(17, 4, 32, device="cuda") * 0.1 + cache = SimpleNamespace( + src=torch.arange(17, device="cuda") % 7, + D_full=torch.eye(16, device="cuda").repeat(17, 1, 1), + ) + seen = {} + original = vp._rotate_mix_scalar_op + + def capture_spy(*args): + u, scalar = original(*args) + # A distinct sentinel proves provenance: a wrong slice has identical + # normal values and would otherwise pass an equality-only assertion. + scalar = scalar + 0.125 + seen["scalar"] = scalar + return u, scalar + + original_alpha = conv._focus_alpha + + def alpha_spy(gate_src): + if capture: + torch.testing.assert_close(gate_src, seen["scalar"], rtol=0, atol=0) + return original_alpha(gate_src) + + with ( + patch.object( + vp, "_rotate_mix_scalar_op", side_effect=capture_spy + ) as scalar_call, + patch.object(vp, "_rotate_mix_op", wraps=vp._rotate_mix_op) as default_call, + patch.object(conv, "_focus_alpha", side_effect=alpha_spy) as alpha_call, + ): + path(x, cache, radial) + assert scalar_call.call_count == int(capture) + assert default_call.call_count == int(not capture) + assert alpha_call.call_count == 1 + + +def _model_params(): + """A small two-block model with the supported scalar-capture layout.""" + return { + "type": "SeZM", + "type_map": ["C"], + "use_compile": False, + "descriptor": { + "type": "SeZM", + "sel": [32], + "rcut": 4.0, + "channels": 32, + "n_focus": 2, + "focus_dim": 32, + "n_radial": 4, + "radial_mlp": [32], + "use_env_seed": True, + "l_schedule": [3, 3, 0], + "mmax": 1, + "so2_layers": 3, + "n_atten_head": 1, + "ffn_neurons": 32, + "ffn_blocks": 1, + "use_amp": False, + "precision": "float32", + "seed": 7, + }, + "fitting_net": {"neuron": [16], "precision": "float32", "seed": 7}, + } + + +@GPU +def test_model_capture_on_off(monkeypatch, strict_environment): + """A real frozen model preserves outputs and executes scalar capture.""" + from deepmd.pt.model.model import ( + get_model, + ) + + model = get_model(_model_params()).to("cuda").eval() + model.requires_grad_(False) + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.2, 0.2, 0.1], [0.3, 1.1, 0.2], [0.1, 0.2, 1.3]]], + device="cuda", + dtype=torch.float64, + ) + atype = torch.zeros(1, 4, device="cuda", dtype=torch.long) + # A test-only ablation leaves the production acceleration policy intact. + with patch.object(vp, "_save_competition_scalar", return_value=False): + reference = model(coord, atype, None) + with patch.object( + vp, "_rotate_mix_scalar_op", wraps=vp._rotate_mix_scalar_op + ) as spy: + actual = model(coord, atype, None) + assert spy.call_count > 0 + assert actual.keys() == reference.keys() + for name in actual: + if isinstance(actual[name], torch.Tensor): + assert torch.isfinite(actual[name]).all() + torch.testing.assert_close( + actual[name], reference[name], rtol=1e-4, atol=1e-5 + ) + + +@pytest.mark.parametrize("backend", ["pt", "pt_expt"]) +@pytest.mark.parametrize( + "level,target", + [ + pytest.param(2, "cuda", marks=GPU), + pytest.param(1, "cuda", marks=GPU), + (2, "cpu"), + ], +) +def test_export_dispatch( + monkeypatch, strict_environment, tmp_path, backend, level, target +): + """Official exporters respect the level, target and storage-lifetime policy. + + Stop at the export boundary, avoiding compiler/toolchain dependencies in + this routing regression. Full AOTI package inference is benchmarked separately. + """ + from deepmd.pt.entrypoints import ( + freeze_pt2, + ) + from deepmd.pt.model.model import ( + get_model, + ) + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + from deepmd.pt_expt.kernels import ( + autotune, + ) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export_impl, + ) + + monkeypatch.setenv("DP_TRITON_INFER", str(level)) + monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None) + monkeypatch.setattr(autotune, "run_autotune", lambda *args: None) + params = _model_params() + model = get_model(params).cpu().eval() + seen = [] + + class ExportBoundary(Exception): + pass + + def inspect_graph(traced, *args, **kwargs): + targets = [str(node.target) for node in traced.graph.nodes] + count = sum("sezm_triton.so2_rotate_mix_scalar." in name for name in targets) + seen.append(count) + assert (count > 0) == (backend == "pt" and level == 2 and target == "cuda") + raise ExportBoundary + + monkeypatch.setattr(torch.export, "export", inspect_graph) + with pytest.raises(ExportBoundary): + if backend == "pt": + checkpoint = tmp_path / "model.pt" + wrapper = ModelWrapper(model, model_params=params) + torch.save({"model": wrapper.state_dict()}, checkpoint) + freeze_pt2.freeze_sezm_to_pt2( + str(checkpoint), + str(tmp_path / "model.pt2"), + device=torch.device(target), + with_comm=False, + ) + else: + _trace_and_export_impl( + {"model": model.serialize(), "model_def_script": params}, + lower_kind="graph", + do_atomic_virial=True, + target_device=torch.device(target), + ) + assert len(seen) == 1 diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py new file mode 100644 index 0000000000..6cc3aaaee0 --- /dev/null +++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Benchmark comparisons must use the same checkpoint and input frames.""" + +from types import ( + SimpleNamespace, +) + +import numpy as np +import pytest + +from .compare_so2_scalar_capture import ( + compare, + configure_capture_route, + validate_identity, +) + + +@pytest.mark.parametrize("route", ["BASE", "CSC5", "TORCH"]) +@pytest.mark.parametrize("available", [False, True]) +@pytest.mark.parametrize("eligible", [False, True]) +def test_capture_route_support(route, available, eligible): + """A pre-feature checkout cannot masquerade as a completed CSC5 run.""" + scalar_path = SimpleNamespace() + if available: + + def original(conv, x): + return eligible + + scalar_path._save_competition_scalar = original + if route == "CSC5" and not available: + with pytest.raises(ValueError, match="CSC5 is unavailable"): + configure_capture_route(route, scalar_path) + return + configure_capture_route(route, scalar_path) + if available: + assert scalar_path._save_competition_scalar(None, None) == ( + eligible and route != "BASE" + ) + if route != "BASE": + assert scalar_path._save_competition_scalar is original + + +@pytest.mark.parametrize("key", ["model_sha256", "input_sha256"]) +@pytest.mark.parametrize("missing", [False, True]) +def test_reject_mismatched_identity(key, missing): + """Reject a changed or absent identity instead of reporting accuracy.""" + result = {"input_sha256": "input"} + manifest = {"model_sha256": "model"} + right_result, right_manifest = dict(result), dict(manifest) + right = right_result if key == "input_sha256" else right_manifest + if missing: + right.pop(key) + else: + right[key] = "different" + with pytest.raises(ValueError, match=key): + validate_identity(result, right_result, manifest, right_manifest) + + +def test_matching_identity(): + """An identical checkpoint and input pair remains comparable.""" + result = {"input_sha256": "input"} + manifest = {"model_sha256": "model"} + validate_identity(result, result, manifest, manifest) + + +def test_compare_checks_cross_repeat_pairs(): + """Equal matching repeats must not hide variability across repeats.""" + scales = {"energy": 1.0, "forces": 10.0, "virial": 100.0, "stress": 1000.0} + arrays = { + f"{repeat}_{frame}_{field}": np.array([(1.0 + repeat) * scales[field]]) + for repeat in range(3) + for frame in range(3) + for field in scales + } + result = compare(arrays, arrays) + for field, scale in scales.items(): + assert not result[field]["passed"] + assert result[field]["max_abs"] == 2.0 * scale + + +@pytest.mark.parametrize("changed_frame", [0, 1, 2]) +def test_compare_uses_reference_and_each_frame(changed_frame): + """Distinct routes must expose a discrepancy confined to any one frame.""" + scales = {"energy": 1.0, "forces": 10.0, "virial": 100.0, "stress": 1000.0} + candidate = { + f"{repeat}_{frame}_{field}": np.array([scale]) + for repeat in range(3) + for frame in range(3) + for field, scale in scales.items() + } + reference = {key: value.copy() for key, value in candidate.items()} + for field, scale in scales.items(): + reference[f"0_{changed_frame}_{field}"] += 2 * scale + result = compare(candidate, reference) + for field, scale in scales.items(): + assert not result[field]["passed"] + assert result[field]["max_abs"] == 2 * scale + assert result[field]["worst"]["frame"] == changed_frame + assert result[field]["worst"]["reference_repeat"] == 0 + assert result[field]["worst"]["index"] == [0] + assert result[field]["worst"]["candidate"] == scale + assert result[field]["worst"]["reference"] == 3 * scale