Skip to content

vibe coded bmm sharder - #3399

Closed
suyoggupta wants to merge 1 commit into
NVIDIA:mainfrom
suyoggupta:user/sg/bmm-sharding
Closed

vibe coded bmm sharder#3399
suyoggupta wants to merge 1 commit into
NVIDIA:mainfrom
suyoggupta:user/sg/bmm-sharding

Conversation

@suyoggupta

Copy link
Copy Markdown
Collaborator

No description provided.

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>

@lucaslie lucaslie left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think it would take to write a pattern matcher for

bmm1 --> bmm2 --> all_reduce

?

@suyoggupta

Copy link
Copy Markdown
Collaborator Author

What do you think it would take to write a pattern matcher for

bmm1 --> bmm2 --> all_reduce

?

We'd need to add all_reduce if we're sharding a dimension that is reduced. The current sharding strategy is to shard the 0th dimension, which happens to be of size num_experts. So the strategy of bmm1->all_gather->bmm2->all_gather works. The first all_gather can be optimized away if all ops between bmm1 and bmm2 are elementwise ops. Does that make sense?

@suyoggupta

Copy link
Copy Markdown
Collaborator Author

open this one instead. nv-auto-deploy#1.

@suyoggupta suyoggupta closed this Apr 10, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants