🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/156868

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

⏳ 1 Pending, 1 Unrelated Failure

As of commit c043984 with merge base 900fba4 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Nice - a few comments:

(1) I think there are a few other places where we rely on has_symints throughout dtensor code, not just that one call site?

(2) we can also kill the has_symint calculation that DTensor does, now that we no longer need it

(3) finally, let's add some tests. I mentioned that you can use the test_dtensor_compile.py tests for inspiration, but in particular:

(3a) we should try to add a test that would have failed without your PR. let me know if you want to brainstorm on ways to exercise this path

(3b) I would also like to generally beef up the amount of dynamic shape testing that we do for DTensor. Doesn't have to be part of this PR. we can talk more about it later - one option is just to manually add more teseting in that file. Another oculd be to hook into existing DTensor OpInfo tests if there are any, let's chat with DTensor folks about it

After clarifying details, I tested several approaches for handling the cache skip. I tested
I tested

1. has_symints -This is the existing approach that just checks top level for symints. This is calculated in the post init of the op info instead of at time of sharding propagation calculator.
2. _are_we_tracing - This replaces the has_symints check with a call to an _are_we_tracing function that already exists. This checks the dispatcher mode and assumes fake tensor dispatch will mean compile globally disables the cache if so.
3. Try/except fallback - This catches a TypeError from the LRUCache attempt and will fallback to the non cached version
4. Hashable conditional - This checks at sharding propagation time if the argument is hashable. This is done deeply, but requires a try/except for unrecognized leaves types as Hashable is not trustable.

Because the change also impacts compiled autograd (and potentially custom backends?), I tested

1. Add
2. Cat (new test that fails with the previous conditional)
3. Compiled autograd of a simple model (existing test that fills with just changing has_symints to _are_we_tracing)
   The following benchmarks are in seconds, and do not include operation time, only the cache logic overhead (including checking the cache).

Add

Cat

Compiled Autograd

Since we would like to prioritize eager runtime (adding compile time is not a big deal), try/except makes the most sense since it is the only time effective method, only adding significant overhead if entering the except clause, that doesn't cause errors. The results are noisy, so in general I expect no to very little additional overhead from this try/except vs the existing conditional statement.

You can find a hacky implementation of the benchmark code here azahed98@e76885f

(1) I think there are a few other places where we rely on has_symints throughout dtensor code, not just that one call site?

@bdhirsh There is one place that uses has_symints (in redistribute), but it recalculates it based on DTensorSpec directly (rather than searching for it), so I left it as is. If I'm understanding correctly, this is fine?

After follow up, the implementation is changed to add a global flag to compiled autograd torch._dynamo.compiled_autograd.in_compiled_autograd_initial_trace. This handles the compiled autograd failure with minimal additional overhead (similar to just _are_we_tracing() above).

@pytorchbot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

@pytorchbot revert -m="This is breaking tests in torchtitan https://github.com/pytorch/torchtitan/actions/runs/16656486277/job/47142656711"

❌ 🤖 pytorchbot command failed:

@pytorchbot revert: error: the following arguments are required: -c/--classification

usage: @pytorchbot revert -m MESSAGE -c
                          {nosignal,ignoredsignal,landrace,weird,ghfirst}

Try @pytorchbot --help for more info.

@pytorchbot revert -m "This is breaking tests in torchtitan https://github.com/pytorch/torchtitan/actions/runs/16656486277/job/47142656711" -c nosignal

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

Reverting PR 156868 failed

Reason: Command git -C /home/runner/work/pytorch/pytorch revert --no-edit f6d138807f138868de0397936e2bee482c1fb987 returned non-zero exit code 1

Auto-merging torch/distributed/_functional_collectives.py
Auto-merging torch/distributed/tensor/_op_schema.py
Auto-merging torch/distributed/tensor/_sharding_prop.py
CONFLICT (content): Merge conflict in torch/distributed/tensor/_sharding_prop.py
error: could not revert f6d138807f1... Always disable ShardingPropagation cache if compiling (#156868)
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git revert --continue".
hint: You can instead skip this commit with "git revert --skip".
hint: To abort and get back to the state before "git revert",
hint: run "git revert --abort".
hint: Disable this message with "git config set advice.mergeConflict false"