[Replicate][Test] tests that pp model grads are the same as single-device model grads

[anshul-si]

Summary: Created a test so that we can verify that a model that has been pipelined + replicated has the same gradients as a reference model. To do this, I mapped the layers and their parameters in each partial model to the original full model and then compared the gradients.

Test Case

1. pytest test/distributed/_composable/test_composability/test_pp_composability.py -k test_replicate_pp_grads

Stack from ghstack (oldest at bottom):

• -> [Replicate][Test] tests that pp model grads are the same as single-device model grads #164880

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

[pytorch-bot]

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Driver update on H100 and A100 instances

⏳ No Failures, 23 Pending

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

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