Skip to content

fix(ATen): honor deterministic algorithms in embedding backward - #152

Open
Hldao wants to merge 2 commits into
MooreThreads:mainfrom
Hldao:fix/embedding-backward-determinism
Open

fix(ATen): honor deterministic algorithms in embedding backward#152
Hldao wants to merge 2 commits into
MooreThreads:mainfrom
Hldao:fix/embedding-backward-determinism

Conversation

@Hldao

@Hldao Hldao commented Aug 18, 2026

Copy link
Copy Markdown

Summary

The short-indices fast path in EmbeddingDenseBwdMUSA accumulates repeated
embedding IDs with atomic additions. The path is selected for
num_indices <= 3072 when scale_grad_by_freq is disabled, even when the user
has enabled deterministic algorithms.

This change keeps the atomic fast path as the default, but bypasses it when
at::globalContext().deterministicAlgorithms() is true. Deterministic mode then
reuses the existing sorted/segmented implementation.

User-visible behavior

  • torch.use_deterministic_algorithms(False): unchanged; eligible calls still
    use the atomic fast path.
  • torch.use_deterministic_algorithms(True): repeated calls with identical
    inputs use the existing deterministic sorted/segmented path.
  • Forward behavior and the mathematical definition of the gradient are
    unchanged.

This follows PyTorch's deterministic-algorithm policy: a nondeterministic
atomic optimization should not be selected when deterministic algorithms are
requested. PyTorch's CUDA short-index dense embedding kernel explicitly
serializes duplicate-ID accumulation instead of using unordered atomic adds.

Regression test

The new test uses 2,063 BF16 indices with repeated IDs, deliberately remaining
below the 3,072-element fast-path threshold. It enables deterministic
algorithms, evaluates aten::embedding_dense_backward eight times, checks
bitwise equality, and restores the caller's original global deterministic
and warn-only state in a finally block.

Device validation

The issue and the proposed routing guard were tested on:

  • device: MTT S4000 48 GB;
  • Python 3.10.8;
  • PyTorch 2.2.0;
  • public source baseline: torch_musa v1.3.0;
  • isolated patched wheel version: torch_musa 1.3.0+3d6a817.

Observed with eight fresh processes per mode:

Mode Unique output hashes
Native atomic path, deterministic mode off 8
Patched deterministic mode 1
Existing sorted-path padding workaround 1

The patched deterministic result exactly matched the sorted-path workaround.
The focused regression test failed on the installed unpatched wheel and passed
on the isolated patched wheel (1 passed).

For this synthetic operator shape, after 25 warm-up calls and 200 synchronized
samples, median latency was 0.9701 ms for the native atomic path and 1.0233 ms
for the patched deterministic path (1.055x). This is an operator microbenchmark,
not an end-to-end training throughput claim. Default-mode routing is unchanged.

Validation boundary

The device validation above was performed against the v1.3-era runtime, not the
current main runtime. The installed vendor wheel reported a private commit
that is not present in the public Git history, so the validation wheel was
built from the public v1.3.0 tag with an ABI-compatible reconstruction using the
installed operator schema and applicable compatibility patches. It validates
the fix mechanism on the stated S4000 environment, but is not a bit-for-bit
reproduction of that private vendor build.

This PR is a static, minimal adaptation of the same guard to current main at
base commit 246d43e3525a9102162cbebfb25c067d634d95a2. Maintainer CI should run
the added test on the current supported MUSA devices and runtime before merge.

Test plan

  • git diff --check
  • Python syntax check for tests/unittest/operator/test_embedding.py
  • Focused regression test on the isolated v1.3 patched wheel: 1 passed
  • Eight-process S4000 reproducer: deterministic mode changed from multiple
    hashes to one hash
  • Current-main MUSA CI / device test (maintainer environment required)

References

Contributor

Prepared by WKK AI R&D, the public-facing brand of
哇咔咔人工智能技术研发有限公司.

@Hldao
Hldao marked this pull request as ready for review August 19, 2026 06:57
Comment thread torch_musa/csrc/aten/ops/musa/Embedding.mu
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants