From c109b5776075c7fb6b542d48c4a2573d0f96ce43 Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 19 Aug 2025 20:36:02 -0700 Subject: [PATCH 1/8] add determinism alerts --- openequivariance/extension/libtorch_tp_jit.cpp | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/extension/libtorch_tp_jit.cpp index dca42fc0..18b8f65c 100644 --- a/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/extension/libtorch_tp_jit.cpp @@ -464,9 +464,11 @@ torch::Tensor jit_conv_forward( check_tensor(rows, {nnz}, k.idx_dtype, "rows"); check_tensor(cols, {nnz}, k.idx_dtype, "cols"); - if (k.deterministic) + if (k.deterministic){ check_tensor(transpose_perm, {nnz}, k.idx_dtype, "transpose perm"); - + } else { + at::globalContext().alertNotDeterministic("OpenEquivariance_conv_atomic_forward"); + } if (k.shared_weights) check_tensor(W, {k.weight_numel}, k.weight_dtype, "W"); else @@ -519,8 +521,11 @@ tuple jit_conv_backward( check_tensor(rows, {nnz}, k.idx_dtype, "rows"); check_tensor(cols, {nnz}, k.idx_dtype, "cols"); - if (k.deterministic) + if (k.deterministic){ check_tensor(transpose_perm, {nnz}, k.idx_dtype, "transpose perm"); + } else { + at::globalContext().alertNotDeterministic("OpenEquivariance_conv_atomic_backward"); + } if (k.shared_weights) check_tensor(W, {k.weight_numel}, k.weight_dtype, "W"); @@ -587,8 +592,11 @@ tuple jit_conv_doubl check_tensor(rows, {nnz}, k.idx_dtype, "rows"); check_tensor(cols, {nnz}, k.idx_dtype, "cols"); - if (k.deterministic) + if (k.deterministic) { check_tensor(transpose_perm, {nnz}, k.idx_dtype, "transpose perm"); + } else { + at::globalContext().alertNotDeterministic("OpenEquivariance_conv_atomic_double_backward"); + } if (k.shared_weights) { check_tensor(W, {k.weight_numel}, k.weight_dtype, "W"); From 29689504a865f60081ca05c1e25e692e17573a2c Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 19 Aug 2025 20:40:04 -0700 Subject: [PATCH 2/8] add determinism tests --- tests/torch_determinism_test.py | 73 +++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 tests/torch_determinism_test.py diff --git a/tests/torch_determinism_test.py b/tests/torch_determinism_test.py new file mode 100644 index 00000000..c8fa83ef --- /dev/null +++ b/tests/torch_determinism_test.py @@ -0,0 +1,73 @@ +import pytest +import torch + +from openequivariance import TPProblem, TensorProductConv + +from e3nn import o3 +from torch_geometric import EdgeIndex + + +@pytest.fixture +def gen(): + return torch.Generator(device="cuda") + + +@pytest.fixture +def edge_index(): + return EdgeIndex( + data=[ + [0, 1, 1, 2], # Receiver + [1, 0, 2, 1], # Sender + ], + sparse_size=(3, 4), + device="cuda", + dtype=torch.long, + ) + + +@pytest.fixture +def tpp(): + X_ir = o3.Irreps("1x2e") + Y_ir = o3.Irreps("1x3e") + Z_ir = o3.Irreps("1x2e") + instructions = [(0, 0, 0, "uvu", True)] + return TPProblem( + X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False + ) + + +@pytest.fixture +def conv_buffers(edge_index, tpp, gen): + X = torch.rand( + edge_index.num_rows, tpp.irreps_in1.dim, device="cuda", generator=gen + ) + Y = torch.rand( + edge_index.num_cols, tpp.irreps_in2.dim, device="cuda", generator=gen + ) + W = torch.rand(edge_index.num_cols, tpp.weight_numel, device="cuda", generator=gen) + return (X, Y, W, edge_index[0], edge_index[1]) + + +@pytest.fixture +def tp_conv(tpp): + return TensorProductConv(tpp, deterministic=False) + + +def test_no_response(tp_conv, conv_buffers): + torch.use_deterministic_algorithms(False) + tp_conv(*conv_buffers) + + +def test_warning(tp_conv, conv_buffers, capfd): + torch.use_deterministic_algorithms(True, warn_only=True) + tp_conv(*conv_buffers) + + captured = capfd.readouterr() + assert "Warning" in captured.err + assert "does not have a deterministic implementation" in captured.err + + +def test_error(tp_conv, conv_buffers): + torch.use_deterministic_algorithms(True, warn_only=False) + with pytest.raises(RuntimeError): + tp_conv(*conv_buffers) From 8bfb0c435f0ea91b05d3d7ee6e4056618929ad11 Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Wed, 8 Oct 2025 13:54:06 -0700 Subject: [PATCH 3/8] nequix tests --- openequivariance/benchmark/problems.py | 10 ++++++++ tests/benchmark.py | 33 +++++++++++++++----------- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/openequivariance/benchmark/problems.py b/openequivariance/benchmark/problems.py index 68822eb9..75d1f47c 100644 --- a/openequivariance/benchmark/problems.py +++ b/openequivariance/benchmark/problems.py @@ -150,6 +150,16 @@ def nequip_problems(): ] ] +# https://github.com/PASSIONLab/OpenEquivariance/discussions/157#discussioncomment-14279135 +def nequix_problems(): + return [ + CTPP( + "128x0e + 64x1o + 32x2e + 32x3o", + "0e + 1o + 2e + 3o", + "128x0e + 64x1o + 32x2e + 32x3o", + "nequix" + ) + ] def e3tools_problems(): return [ diff --git a/tests/benchmark.py b/tests/benchmark.py index ab005cdf..82a5e640 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -46,6 +46,7 @@ diffdock_problems, mace_problems, nequip_problems, + nequix_problems, ) from torch._functorch import config @@ -91,11 +92,14 @@ def benchmark_uvu(params): - float64_problems = mace_problems() + nequip_problems() + def get_problems(): + return mace_problems() + nequip_problems() + nequix_problems() + + float64_problems = get_problems() for problem in float64_problems: problem.irrep_dtype = np.float64 problem.weight_dtype = np.float64 - problems = mace_problems() + nequip_problems() + float64_problems + problems = get_problems() + float64_problems implementations = [implementation_map_tp[impl] for impl in params.implementations] directions = params.directions @@ -280,18 +284,19 @@ def benchmark_convolution(params): graphs = download_graphs(params, filenames) if not params.disable_bench: - configs = [ - ChannelwiseTPP( - "128x0e+128x1o+128x2e", - "1x0e+1x1o+1x2e+1x3o", - "128x0e+128x1o+128x2e+128x3o", - ), - ChannelwiseTPP( - "128x0e+128x1o+128x2e", - "1x0e+1x1o+1x2e+1x3o", - "128x0e+128x1o+128x2e+128x3o", - ), - ] # MACE-large + configs = nequix_problems() + nequix_problems() + # [ + # ChannelwiseTPP( + # "128x0e+128x1o+128x2e", + # "1x0e+1x1o+1x2e+1x3o", + # "128x0e+128x1o+128x2e+128x3o", + # ), + # ChannelwiseTPP( + # "128x0e+128x1o+128x2e", + # "1x0e+1x1o+1x2e+1x3o", + # "128x0e+128x1o+128x2e+128x3o", + # ), + # ] # MACE-large configs[1].irrep_dtype = np.float64 configs[1].weight_dtype = np.float64 From c6a6e5b91d6d53daa1e72327c87d924fc1ba126e Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 14 Oct 2025 21:40:20 -0700 Subject: [PATCH 4/8] add problems --- openequivariance/benchmark/problems.py | 45 +++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/openequivariance/benchmark/problems.py b/openequivariance/benchmark/problems.py index 75d1f47c..07b1d823 100644 --- a/openequivariance/benchmark/problems.py +++ b/openequivariance/benchmark/problems.py @@ -150,17 +150,52 @@ def nequip_problems(): ] ] -# https://github.com/PASSIONLab/OpenEquivariance/discussions/157#discussioncomment-14279135 +# https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml def nequix_problems(): return [ CTPP( - "128x0e + 64x1o + 32x2e + 32x3o", - "0e + 1o + 2e + 3o", - "128x0e + 64x1o + 32x2e + 32x3o", - "nequix" + "89x0e", + "1x0e+1x1o+1x2e+1x3o", + "89x0e+89x1o+89x2e+89x3o", + "nequix-mp-1-first_layer" + ), + CTPP( + "128x0e+64x1o+32x2e+32x3o", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+128x1o+128x2e+128x3o+64x1o+64x0e+64x2e+64x1o+64x3o+64x2e+32x2e+32x1o+32x3o+32x0e+32x2e+32x1o+32x3o+32x3o+32x2e+32x1o+32x3o+32x0e+32x2e", + "nequix-mp-1-main_layers" + ), + CTPP( + "128x0e+64x1o+32x2e+32x3o", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+64x0e+32x0e+32x0e", + "nequix-mp-1-last_layer" ) ] +# https://github.com/MDIL-SNU/SevenNet/tree/main/sevenn/pretrained_potentials/SevenNet_l3i5 +def seven_net_problems(): + return [ + CTPP( + "128x0e", + "1x0e+1x1e+1x2e+1x3e", + "128x0e+128x1e+128x2e+128x3e", + "SevenNet_l3i5-first-layer" + ), + CTPP( + "128x0e+64x1e+32x2e+32x3e", + "1x0e+1x1e+1x2e+1x3e", + "128x0e+64x0e+32x0e+32x0e+128x1e+64x1e+64x1e+64x1e+32x1e+32x1e+32x1e+32x1e+32x1e+128x2e+64x2e+64x2e+64x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+128x3e+64x3e+64x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e", + "SevenNet_l3i5-main-layers" + ), + CTPP( + "128x0e+64x1e+32x2e+32x3e", + "1x0e+1x1e+1x2e+1x3e", + "128x0e+64x0e+32x0e+32x0e", + "SevenNet_l3i5-last-layer" + ), + ] + def e3tools_problems(): return [ FCTPP(in1, in2, out, label=label, shared_weights=sw, internal_weights=iw) From 9ff6d17acae3d7c25886174258915b91d64b6cf5 Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 14 Oct 2025 21:41:28 -0700 Subject: [PATCH 5/8] format --- openequivariance/benchmark/problems.py | 25 ++++++++++++++----------- tests/benchmark.py | 4 ++-- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/openequivariance/benchmark/problems.py b/openequivariance/benchmark/problems.py index 07b1d823..dd536e3a 100644 --- a/openequivariance/benchmark/problems.py +++ b/openequivariance/benchmark/problems.py @@ -150,6 +150,7 @@ def nequip_problems(): ] ] + # https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml def nequix_problems(): return [ @@ -157,45 +158,47 @@ def nequix_problems(): "89x0e", "1x0e+1x1o+1x2e+1x3o", "89x0e+89x1o+89x2e+89x3o", - "nequix-mp-1-first_layer" + "nequix-mp-1-first_layer", ), CTPP( "128x0e+64x1o+32x2e+32x3o", "1x0e+1x1o+1x2e+1x3o", "128x0e+128x1o+128x2e+128x3o+64x1o+64x0e+64x2e+64x1o+64x3o+64x2e+32x2e+32x1o+32x3o+32x0e+32x2e+32x1o+32x3o+32x3o+32x2e+32x1o+32x3o+32x0e+32x2e", - "nequix-mp-1-main_layers" + "nequix-mp-1-main_layers", ), CTPP( "128x0e+64x1o+32x2e+32x3o", "1x0e+1x1o+1x2e+1x3o", "128x0e+64x0e+32x0e+32x0e", - "nequix-mp-1-last_layer" - ) + "nequix-mp-1-last_layer", + ), ] -# https://github.com/MDIL-SNU/SevenNet/tree/main/sevenn/pretrained_potentials/SevenNet_l3i5 + +# https://github.com/MDIL-SNU/SevenNet/tree/main/sevenn/pretrained_potentials/SevenNet_l3i5 def seven_net_problems(): return [ CTPP( "128x0e", "1x0e+1x1e+1x2e+1x3e", "128x0e+128x1e+128x2e+128x3e", - "SevenNet_l3i5-first-layer" - ), + "SevenNet_l3i5-first-layer", + ), CTPP( "128x0e+64x1e+32x2e+32x3e", "1x0e+1x1e+1x2e+1x3e", "128x0e+64x0e+32x0e+32x0e+128x1e+64x1e+64x1e+64x1e+32x1e+32x1e+32x1e+32x1e+32x1e+128x2e+64x2e+64x2e+64x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+128x3e+64x3e+64x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e", - "SevenNet_l3i5-main-layers" - ), + "SevenNet_l3i5-main-layers", + ), CTPP( "128x0e+64x1e+32x2e+32x3e", "1x0e+1x1e+1x2e+1x3e", "128x0e+64x0e+32x0e+32x0e", - "SevenNet_l3i5-last-layer" - ), + "SevenNet_l3i5-last-layer", + ), ] + def e3tools_problems(): return [ FCTPP(in1, in2, out, label=label, shared_weights=sw, internal_weights=iw) diff --git a/tests/benchmark.py b/tests/benchmark.py index 82a5e640..9563d2d8 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -93,8 +93,8 @@ def benchmark_uvu(params): def get_problems(): - return mace_problems() + nequip_problems() + nequix_problems() - + return mace_problems() + nequip_problems() + nequix_problems() + float64_problems = get_problems() for problem in float64_problems: problem.irrep_dtype = np.float64 From 5d5b8b24141b809b4927f46d5580938558b6c3b6 Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 14 Oct 2025 22:29:04 -0700 Subject: [PATCH 6/8] benchmark nequix and seven-net problems --- tests/benchmark.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/tests/benchmark.py b/tests/benchmark.py index 9563d2d8..fa4d9f4f 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -47,6 +47,7 @@ mace_problems, nequip_problems, nequix_problems, + seven_net_problems, ) from torch._functorch import config @@ -93,7 +94,7 @@ def benchmark_uvu(params): def get_problems(): - return mace_problems() + nequip_problems() + nequix_problems() + return mace_problems() + nequip_problems() + nequix_problems() + seven_net_problems() float64_problems = get_problems() for problem in float64_problems: @@ -284,22 +285,14 @@ def benchmark_convolution(params): graphs = download_graphs(params, filenames) if not params.disable_bench: - configs = nequix_problems() + nequix_problems() - # [ - # ChannelwiseTPP( - # "128x0e+128x1o+128x2e", - # "1x0e+1x1o+1x2e+1x3o", - # "128x0e+128x1o+128x2e+128x3o", - # ), - # ChannelwiseTPP( - # "128x0e+128x1o+128x2e", - # "1x0e+1x1o+1x2e+1x3o", - # "128x0e+128x1o+128x2e+128x3o", - # ), - # ] # MACE-large - - configs[1].irrep_dtype = np.float64 - configs[1].weight_dtype = np.float64 + def get_problems(): + return mace_problems() + nequip_problems() + nequix_problems() + seven_net_problems() + + float64_problems = get_problems() + for problem in float64_problems: + problem.irrep_dtype = np.float64 + problem.weight_dtype = np.float64 + configs = get_problems() + float64_problems bench = ConvBenchmarkSuite(configs, test_name="convolution") From d7af1b0bb1ff068605e33a3bed8f738857870ee2 Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Sun, 9 Nov 2025 21:31:09 -0800 Subject: [PATCH 7/8] format --- tests/benchmark.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/benchmark.py b/tests/benchmark.py index fa4d9f4f..aab38b31 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -94,7 +94,12 @@ def benchmark_uvu(params): def get_problems(): - return mace_problems() + nequip_problems() + nequix_problems() + seven_net_problems() + return ( + mace_problems() + + nequip_problems() + + nequix_problems() + + seven_net_problems() + ) float64_problems = get_problems() for problem in float64_problems: @@ -285,8 +290,14 @@ def benchmark_convolution(params): graphs = download_graphs(params, filenames) if not params.disable_bench: + def get_problems(): - return mace_problems() + nequip_problems() + nequix_problems() + seven_net_problems() + return ( + mace_problems() + + nequip_problems() + + nequix_problems() + + seven_net_problems() + ) float64_problems = get_problems() for problem in float64_problems: From dd7c6c998824ba86977246c78d3e42fee3fcf45e Mon Sep 17 00:00:00 2001 From: Austin Glover Date: Tue, 11 Nov 2025 12:20:44 -0800 Subject: [PATCH 8/8] remove changes to benchmark.py --- tests/benchmark.py | 43 +++++++++++++++++-------------------------- 1 file changed, 17 insertions(+), 26 deletions(-) diff --git a/tests/benchmark.py b/tests/benchmark.py index aab38b31..ab005cdf 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -46,8 +46,6 @@ diffdock_problems, mace_problems, nequip_problems, - nequix_problems, - seven_net_problems, ) from torch._functorch import config @@ -93,19 +91,11 @@ def benchmark_uvu(params): - def get_problems(): - return ( - mace_problems() - + nequip_problems() - + nequix_problems() - + seven_net_problems() - ) - - float64_problems = get_problems() + float64_problems = mace_problems() + nequip_problems() for problem in float64_problems: problem.irrep_dtype = np.float64 problem.weight_dtype = np.float64 - problems = get_problems() + float64_problems + problems = mace_problems() + nequip_problems() + float64_problems implementations = [implementation_map_tp[impl] for impl in params.implementations] directions = params.directions @@ -290,20 +280,21 @@ def benchmark_convolution(params): graphs = download_graphs(params, filenames) if not params.disable_bench: - - def get_problems(): - return ( - mace_problems() - + nequip_problems() - + nequix_problems() - + seven_net_problems() - ) - - float64_problems = get_problems() - for problem in float64_problems: - problem.irrep_dtype = np.float64 - problem.weight_dtype = np.float64 - configs = get_problems() + float64_problems + configs = [ + ChannelwiseTPP( + "128x0e+128x1o+128x2e", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+128x1o+128x2e+128x3o", + ), + ChannelwiseTPP( + "128x0e+128x1o+128x2e", + "1x0e+1x1o+1x2e+1x3o", + "128x0e+128x1o+128x2e+128x3o", + ), + ] # MACE-large + + configs[1].irrep_dtype = np.float64 + configs[1].weight_dtype = np.float64 bench = ConvBenchmarkSuite(configs, test_name="convolution")