{{ message }}
Guard tensors against the current device index under compile_on_one_rank - #196150
Draft
aorenste wants to merge 1 commit into
Draft
Guard tensors against the current device index under compile_on_one_rank#196150aorenste wants to merge 1 commit into
aorenste wants to merge 1 commit into
Conversation
[ghstack-poisoned]
This was referenced Sep 6, 2026
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.

Stack from ghstack (oldest at bottom):
A TENSOR_MATCH guard records the device as two independent pieces: the type
rides in the DispatchKeySet, and the index is a separate scalar. Only the index
is rank-specific, and under compile_on_one_rank it carries no information at
all -- CooR enforces a single-accelerator invariant while tracing, rejecting any
operand or tensor on a non-current accelerator device. Each rank therefore bakes
its own index into an otherwise identical guard (device=0 on rank 0, device=7 on
rank 7), so no rank can reuse another's compiled artifact.
Compare the index against the current accelerator device instead of the one
recorded when the guard was built. This is a relaxation, not a removal: a tensor
on some other device still fails the guard. That matters because the invariant
above is enforced when tracing, leaving the guard as the only thing watching at
runtime. The device type is untouched -- cpu and cuda tensors coexist freely in
one process and that distinction lives in the dispatch key -- and cpu tensors
are skipped entirely, their device already being portable across ranks.
Gated on compile_on_one_rank: outside it several accelerator devices can
legitimately be live at once, so there the index stays pinned.
The guard string renders as device=current rather than an index, because that
text is what gets serialized into a precompile artifact; a correct dynamic check
paired with a rank-specific string would still not be portable.
Measured on a 16-GPU CooR job, these guards were the largest remaining source of
cross-rank divergence: 2744 device=0 entries on rank 0 against 2744 device=7 on
rank 15, in guard sets that were otherwise identical.