fix(ATen): honor deterministic algorithms in embedding backward

tests/unittest/operator/test_embedding.py

@@ -213,6 +213,72 @@ def test_embedding_bwd(dtype, scale_grad_by_freq, input_shape, fixed_indices):
             assert not comparator(scaled_result.float(), unscaled_result.float())
 
 
+@pytest.mark.skipif(
+    testing.get_musa_arch() < 22,
+    reason="bf16 is not supported on arch older than qy2",
+)
+@testing.test_on_nonzero_card_if_multiple_musa_device(1)
+def test_embedding_dense_backward_deterministic_algorithms():
+    indices = torch.tensor(
+        [0] * 133 + [1 + (position % 156) for position in range(2063 - 133)],
+        dtype=torch.int64,
+        device="musa",
+    )
+    positions = torch.arange(2063 * 128, dtype=torch.float32).reshape(2063, 128)
+    grad = (
+        torch.sin(positions * 0.00037) * 0.125
+        + torch.cos(positions * 0.00011) * 0.0625
+    ).to(device="musa", dtype=torch.bfloat16)
+
+    deterministic_algorithms_enabled = (
+        torch.are_deterministic_algorithms_enabled()
+    )
+    deterministic_algorithms_warn_only_enabled = (
+        torch.is_deterministic_algorithms_warn_only_enabled()
+    )
+    try:
+        torch.use_deterministic_algorithms(True)
+        outputs = [
+            torch.ops.aten.embedding_dense_backward(
+                grad, indices, 512, -1, False
+            ).cpu()
+            for _ in range(8)
+        ]
+    finally:
+        torch.use_deterministic_algorithms(
+            deterministic_algorithms_enabled,
+            warn_only=deterministic_algorithms_warn_only_enabled,
+        )
+
+    assert all(torch.equal(outputs[0], output) for output in outputs[1:])
+
+
+@testing.test_on_nonzero_card_if_multiple_musa_device(1)
+def test_embedding_dense_backward_deterministic_algorithms_empty_indices():
+    indices = torch.empty((0,), dtype=torch.int64, device="musa")
+    grad = torch.empty((0, 128), dtype=torch.float32, device="musa")
+
+    deterministic_algorithms_enabled = (
+        torch.are_deterministic_algorithms_enabled()
+    )
+    deterministic_algorithms_warn_only_enabled = (
+        torch.is_deterministic_algorithms_warn_only_enabled()
+    )
+    try:
+        torch.use_deterministic_algorithms(True)
+        output = torch.ops.aten.embedding_dense_backward(
+            grad, indices, 512, -1, False
+        )
+    finally:
+        torch.use_deterministic_algorithms(
+            deterministic_algorithms_enabled,
+            warn_only=deterministic_algorithms_warn_only_enabled,
+        )
+
+    assert output.shape == (512, 128)
+    assert torch.count_nonzero(output).item() == 0
+
+
 float_dtypes = [torch.float32]
 
 

torch_musa/csrc/aten/ops/musa/Embedding.mu

@@ -1,4 +1,5 @@
 #include <ATen/ATen.h>
+#include <ATen/Context.h>
 #include <ATen/Dispatch.h>
 #include <ATen/ceil_div.h>
 #include <ATen/core/Tensor.h>
@@ -72,7 +73,8 @@ Tensor EmbeddingDenseBwdMUSA(
 
   // be careful for setting this value, there may be
   // precision and efficiency drops when value gets larger.
-  if (num_indices <= 3072 && !scale_grad_by_freq) {
+  if (num_indices <= 3072 && !scale_grad_by_freq &&
+      !at::globalContext().deterministicAlgorithms()) {
     Tensor grad_weight = at::zeros(
         {num_weights, grad_output.size(-1)},
         grad_output.options().memory_format(at::MemoryFormat::Contiguous));

torch_musa/csrc/aten/ops/musa/EmbeddingBackwardKernel.mu

@@ -316,6 +316,10 @@ Tensor EmbeddingBackwardMUSAKernel(
   const ptrdiff_t numel = sorted_indices.numel();
   Tensor grad_weight = at::zeros({num_weights, grad.size(-1)}, grad.options());
 
+  if (C10_UNLIKELY(numel == 0)) {
+    return grad_weight;
+  }
+
   int tbl_h = grad_weight.size(0);
   int tbl_w = grad_weight.size(1);