tp rough draft

[Quentin-Anthony]

@joellidin -- Draft TP impl with some cleanups. Some outstanding issues:

1. Need to figure out the outer step logic
2. My updates to prepare_gradient_dict break non-TP and need some improvements to make it bulletproof
3. There's a new grad dtensor gather that we should probably remove if TP is being used
4. The new owned_params needs reworked for TP
5. We need error feedback padding, I think.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch torchtitan-tp

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}