Guard tensors against the current device index under compile_on_one_rank by aorenste · Pull Request #196150 · pytorch/pytorch · GitHub
Skip to content

Guard tensors against the current device index under compile_on_one_rank - #196150

Draft
aorenste wants to merge 1 commit into
gh/aorenste/379/basefrom
gh/aorenste/379/head
Draft

Guard tensors against the current device index under compile_on_one_rank#196150
aorenste wants to merge 1 commit into
gh/aorenste/379/basefrom
gh/aorenste/379/head

Conversation

@aorenste

@aorenste aorenste commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

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.

[ghstack-poisoned]
@pytorch-bot pytorch-bot Bot added ciflow/dtensor Run DTensor specific tests ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests module: dynamo labels Sep 6, 2026
@pytorch-bot

pytorch-bot Bot commented Sep 6, 2026

Copy link
Copy Markdown

@pytorch-bot

pytorch-bot Bot commented Sep 6, 2026

Copy link
Copy Markdown

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/dtensor Run DTensor specific tests ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests module: dynamo

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant