Fix XPUGraph falling back to eager on models with embeddings - #4547
Open
jemitche1 wants to merge 14 commits into
Open
Fix XPUGraph falling back to eager on models with embeddings#4547jemitche1 wants to merge 14 commits into
jemitche1 wants to merge 14 commits into
Conversation
Signed-off-by: Jerome Mitchell <jerome.mitchell@intel.com>
torch-xpu-ops implements aten.embedding_dense_backward with
embedding_backward_deterministic_kernel, which calls Tensor::item<long>()
to size an allocation. That is a host-blocking read of device data, so
SYCL command-graph recording rejects it with "wait method cannot be used
for an event associated with a command graph".
Excluding the node is not an option. is_xpugraph_compatible is a
whole-graph gate, so one unsafe node disables capture entirely rather
than capturing around it -- and the existing _local_scalar_dense.default
exclusion cannot catch this anyway, because the .item() lives inside a
C++ kernel with no separate FX node to exclude.
Rewrite the node instead, into primitives with static shapes and no host
sync:
grad_weight = zeros(num_weights, D).index_add(0, indices, grad)
Nodes with padding_idx >= 0 or scale_grad_by_freq=True are left alone and
logged, since neither is reproduced by this substitution.
Tradeoff: index_add accumulates with atomics, so embedding-gradient
accumulation order is no longer deterministic and embedding grads are not
bitwise reproducible run to run, where the kernel this replaces was
specifically the deterministic one. A deterministic alternative would
require making embedding_backward_deterministic_kernel size its
allocations from an upper bound instead of via .item(), which is a
torch-xpu-ops change.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
jemitche1
requested review from
IvanKobzarev,
SherlockNoMad,
aditvenk,
sanketpurandare and
tianyu-l
as code owners
September 9, 2026 18:20
meta-cla
Bot
added
the
CLA Signed
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
1 participant
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.
This PR rewrites the embedding backward pass in the traced graph so that XPUGraph capture in graph_trainer no longer falls back to eager; 0the kernel behind that operation reads a value off the device to size an allocation, which is illegal mid-recording, and capture is all-or-nothing. A new pass runs inside xpugraph_pass just before the compatibility gate and replaces each aten.embedding_dense_backward node with zeros(num_weights, D) followed by index_add(0, indices, grad), which computes the same gradient using only shapes known up front and no reads back to the host