Record redistribute_local_tensor in DebugMode

[SherlockNoMad]

Explicit redistribute_local_tensor API call could also results in communication, record it!

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim @dcci

[pytorch-bot]

🔗 Helpful Links

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

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8b1ee97 with merge base 8d81564 ():

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

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

• inductor / inductor-cpu-test / test (dynamic_cpu_inductor_torchbench, 2, 2, linux.8xlarge.amx) (gh) (trunk failure)
  torch_multimodal_clip

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

[zpcore]

Nice, I was going to implement capturing the transforminfo but you already did it! Is it possible we have a flag in debug_mode to control this?

Overall LGTM, may need to update the test file.

[zpcore]

By the way, I hope in DebugMode we can get the list of transforminfos triggered for each DTensor redistribution besides the detailed collective ops carried out under the hood. Is it possible? Thanks!

[SherlockNoMad]

@zpcore

I tried

  redistribute_input(t: f32[3, 8], [S(0)] -> [R], ([_TransformInfo(mesh_dim=0, src_dst_placements=(Shard(dim=0), Replicate()), logical_shape=[9, 8])],))

looks a bit too verbose? anyway to simplify?

[ezyang]

Put this in a helper plz

[SherlockNoMad]

done.

get_active_debug_mode

[SherlockNoMad]

@pytorchbot merge

[pytorchmergebot]

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

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• Lint / lintrunner-noclang / linux-job

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[SherlockNoMad]

@pytorchbot merge

[pytorchmergebot]

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

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (distributed, 3, 3, linux.g4dn.12xlarge.nvidia.gpu)

Details for Dev Infra team

Raised by workflow job

[SherlockNoMad]

@pytorchbot merge

[pytorchmergebot]

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

Advanced Debugging

Check the merge workflow status
here

[zpcore]

@zpcore

I tried

  redistribute_input(t: f32[3, 8], [S(0)] -> [R], ([_TransformInfo(mesh_dim=0, src_dst_placements=(Shard(dim=0), Replicate()), logical_shape=[9, 8])],))

looks a bit too verbose? anyway to simplify?

That's good enough, I will follow up to update the str represestation of the transforminfo. Thanks!