[FSDP2] Showed 2D MLP with colwise + colwise sharding

[awgu]

Stack from ghstack (oldest at bottom):

• -> [FSDP2] Showed 2D MLP with colwise + colwise sharding #126073
• [FSDP2] Factored out MLPStack to de-dup code #126070
• [FSDP2] Used CommDebugMode in grad acc test #126067

What is the point of this?

• FSDP2 currently only supports rowwise sharding (i.e. dim-0 sharding). This means that when combined with tensor parallelism, some linear weights are rowwise sharded twice, meaning that one matrix dim is sharded both by TP and FSDP while the other matrix dim is not sharded.
• For possible performance reasons (e.g. quantization block sizes), we may prefer to change those linear weights to have FSDP instead do colwise sharding (i.e. dim-1 sharding). However, implementing that in FSDP2 introduces some complexity. It may be simpler to transpose the weight and change TP to shard colwise. (More verbosely, instead of using TP rowwise sharding and FSDP colwise sharding, we transpose the weight, use TP colwise sharding, and FSDP rowwise sharding.)

cc @XilunWu @H-Huang @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @penguinwu @tianyu-l @yf225 @chauhang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 62ef78a with merge base afda668 ():
💚 Looks good so far! There are no failures yet. 💚

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

[awgu]

cc: @wanchaol @yifuwang curious to get your thoughts here

In my naive understanding, there should not be any noticeable performance implication.

[wanchaol]

For ColwiseParallel, iirc the actual tensor sharding dim would be 0 (this is because there was a transpose before, colwise parallel suppose to shard on 1, but because of the tranpose the weight sharding is on 0).

Also the desired input layouts of ColwiseParallel would be replicated, and rowwise parallel desired input layout would be Shard(-1), so there might be some runtime difference though.

If we want to achieve something that you want: TP shard on dim1, while FSDP2 shard on dim0, I think we can start directly from DTensor APIs (i.e. we copy over Colwise/RowwiseParallel and make adjustments), if this works out we can think of how to expose this option to the default Colwise/RowwiseParallel

[yifuwang]

meaning that one matrix dim is sharded both by TP and FSDP while the other matrix dim is not sharded.

Would it make sense to just support this? Curious if this feels natural to folks (for me it kinda is). If so, how difficult would this be (if possible)?

[awgu]

@yifuwang That is the current way 2D works for FSDP2 (so it is supported, just missing the strided sharding placement that is being worked on, which only affects resharding).

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.