[DTensor] Register replication strategy for a few upsampling interpolate ops

[wz337]

Stack from ghstack (oldest at bottom):

• -> [DTensor] Register replication strategy for a few upsampling interpolate ops #137201

To unblock Llama 3.2 vision's use case to resize positional embeddings for fine-tuning. Context in workplace post.

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wconstab @d4l3k @c-p-i-o @tianyu-l

[pytorch-bot]

🔗 Helpful Links

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

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

✅ No Failures

As of commit 4abdf2a with merge base 8962610 ():
💚 Looks good so far! There are no failures yet. 💚

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

[XilunWu]

Replicate strategy should be fine for now.

[wz337]

Replicate strategy should be fine for now.

Ye. I look at the ops and it doesn't seem we can improve sharding strategy for these ones so it would always require redistributing to replicate.

[awgu]

How would the replicated DTensor get converted back to sharded in the state dict load flow?

[wz337]

How would the replicated DTensor get converted back to sharded in the state dict load flow?

Ah. Thanks for raising it. We won't be able to define the layout for the output, since it is determined by the next op.

In order to shard it back, users would have to do a dtensor.redistribute(placements=[Shard(0)]). This way there is no communication compared with using distribute_tensor().

[wz337]

@pytorchmergebot 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