fix(ATen): honor deterministic algorithms in embedding backward - #152
Open
Hldao wants to merge 2 commits into
Open
fix(ATen): honor deterministic algorithms in embedding backward#152Hldao wants to merge 2 commits into
Hldao wants to merge 2 commits into
Conversation
Hldao
marked this pull request as ready for review
August 19, 2026 06:57
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The short-indices fast path in
EmbeddingDenseBwdMUSAaccumulates repeatedembedding IDs with atomic additions. The path is selected for
num_indices <= 3072whenscale_grad_by_freqis disabled, even when the userhas enabled deterministic algorithms.
This change keeps the atomic fast path as the default, but bypasses it when
at::globalContext().deterministicAlgorithms()is true. Deterministic mode thenreuses the existing sorted/segmented implementation.
User-visible behavior
torch.use_deterministic_algorithms(False): unchanged; eligible calls stilluse the atomic fast path.
torch.use_deterministic_algorithms(True): repeated calls with identicalinputs use the existing deterministic sorted/segmented path.
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_backwardeight times, checksbitwise equality, and restores the caller's original global deterministic
and warn-only state in a
finallyblock.Device validation
The issue and the proposed routing guard were tested on:
torch_musa 1.3.0+3d6a817.Observed with eight fresh processes per mode:
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
mainruntime. The installed vendor wheel reported a private committhat 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
mainatbase commit
246d43e3525a9102162cbebfb25c067d634d95a2. Maintainer CI should runthe added test on the current supported MUSA devices and runtime before merge.
Test plan
git diff --checktests/unittest/operator/test_embedding.py1 passedhashes to one hash
References
https://github.com/Hldao/musa-embedding-determinism
https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/EmbeddingBackwardKernel.cu
Contributor
Prepared by WKK AI R&D, the public-facing brand of
哇咔咔人工智能技术研发有限公司.