diff --git a/tests/unittest/operator/test_embedding.py b/tests/unittest/operator/test_embedding.py index da91531..c511355 100644 --- a/tests/unittest/operator/test_embedding.py +++ b/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] diff --git a/torch_musa/csrc/aten/ops/musa/Embedding.mu b/torch_musa/csrc/aten/ops/musa/Embedding.mu index 7e534b3..c5972e7 100644 --- a/torch_musa/csrc/aten/ops/musa/Embedding.mu +++ b/torch_musa/csrc/aten/ops/musa/Embedding.mu @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -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)); diff --git a/torch_musa/csrc/aten/ops/musa/EmbeddingBackwardKernel.mu b/torch_musa/csrc/aten/ops/musa/EmbeddingBackwardKernel.mu index 362c0f2..8031c9c 100644 --- a/torch_musa/csrc/aten/ops/musa/EmbeddingBackwardKernel.mu +++ b/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);