From 4becbe81ef914fff10ab4e4750d61fe60bde12af Mon Sep 17 00:00:00 2001 From: Tingyu Wang Date: Wed, 20 Mar 2024 16:02:08 -0400 Subject: [PATCH 1/2] Always scatter in FP32 in TPConv (#4260) Always scatter in FP32 in TPConv due to poor performance of torch's native scatter_add in bfloat16. Authors: - Tingyu Wang (https://github.com/tingyu66) Approvers: - Alex Barghi (https://github.com/alexbarghi-nv) URL: https://github.com/rapidsai/cugraph/pull/4260 --- .../cugraph_equivariant/nn/tensor_product_conv.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py index 5120a23180d..af1d0efa76c 100644 --- a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py +++ b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py @@ -251,7 +251,10 @@ def forward( if edge_envelope is not None: out = out * edge_envelope.view(-1, 1) - out = scatter_reduce(out, dst, dim=0, dim_size=num_dst_nodes, reduce=reduce) + dtype = out.dtype + out = scatter_reduce( + out.float(), dst, dim=0, dim_size=num_dst_nodes, reduce=reduce + ).to(dtype) if self.batch_norm: out = self.batch_norm(out) From 87061152bfad49f4a5d73288eed399b773218a0a Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Thu, 21 Mar 2024 09:12:22 -0400 Subject: [PATCH 2/2] Update pre-commit-hooks to v0.0.3 (#4264) This fixes an issue with how the `verify-copyright` hook handles multiple merge bases. Authors: - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - Don Acosta (https://github.com/acostadon) URL: https://github.com/rapidsai/cugraph/pull/4264 --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 542e9cacb77..a314b8c7185 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -52,7 +52,7 @@ repos: - id: rapids-dependency-file-generator args: ["--clean"] - repo: https://github.com/rapidsai/pre-commit-hooks - rev: v0.0.1 + rev: v0.0.3 hooks: - id: verify-copyright files: |