DTensor: use memory_format in the hash for all aten ops that use that arg (e.g. aten.clone)

[bdhirsh]

This fixes an internal DTensor enablement bug (I don't have an OSS issue for it)

I finally root-caused this as follows:

(1) we were fakefying a DTensor graph input, that was an autograd non-leaf (it had a grad_fn)

(2) that caused it do go through this clone() call during fakeification: https://github.com/pytorch/pytorch/blob/main/torch/_subclasses/meta_utils.py#L549

(3) clone(torch.preserve_format) is supposed to return another DTensor with the same strides as the input, but I noticed we were returning a DTensor with contiguous strides incorrectly.

(4) It turns out that DTensor was hashing on the sharding strategy for aten.clone, regardless of the memory_format kwarg that was passed in.

I could have manually updated the clone sharding strategy registration to take memory_format into account. But instead, I figured that every aten op with a sharding strategy needs to handle the memory_format kwarg specially - so I tried to generically force DTensor to consider all ATen ops that take a memory_format kwarg during hashing.

Stack from ghstack (oldest at bottom):

• [not ready yet] potentially fix set_() path of fakeifying for subclasses #119948
• [not ready yet] guard on tensor subclass types before running inner tensor guards #119067
• inductor: fix for functional_collectives.wait() followed by view() #118802
• DTensor: when forcing tangents to be correct, force syncing collectives at graph boundaries #118671
• AOTDispatch: allow subclasses to correct when we guess metadata of tangents incorrectly #118670
• AOTAutograd: ensure traced tangent subclass metadata takes non-contiguous outputs into account #118669
• dynamo: support placement kwargs for DTensor.to_local() #119947
• dynamo: handle DTensor.device_mesh.device_type #118803
• DTensor + dynamo: fix is_shard/replicate always inlining to False #118668
• -> DTensor: use memory_format in the hash for all aten ops that use that arg (e.g. aten.clone) #118667
• fix multiple-fake-modes bug with compile + subclasses #118191
• DTensor: make tensor_flatten more compatible for dynamo getattr #118209
• dynamo: support attribute access on tensor subclasses without sources #117666
• dynamo: respect autograd.Function + multiple save_for_backward calls #117667

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 3c91fd8 with merge base 86dedeb ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / linux-focal-cuda12.1-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  Build completed, 1 test FAILED, 46 total actions

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / rocm6.0-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2, unstable) (gh)
  Process completed with exit code 1.

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

[bdhirsh]

cc @wanchaol, let me know who else it would be good to get a review from from the distributed side

[wanchaol]

lgtm! thanks for fixing the bug and make it applicable to other cases!