Skip to content

Conversation

copybara-service[bot]
Copy link

Always propagate shardings along the batch dimension in InferConvolutionShardingFromOperands.

Before this cl, if the kernel is large, we replicate the output, even along the batch dimension. This cl enables the sharding propagation between lhs and output along the batch dimension such that we can support data parallelism.

Add a test case, which fails before this cl and succeeds with this cl.

@copybara-service copybara-service bot force-pushed the test_614049496 branch 2 times, most recently from 41e0214 to ba2b116 Compare March 11, 2024 17:45
…tionShardingFromOperands`.

Before this cl, if the kernel is large, we replicate the output, even along the batch dimension. This cl enables the sharding propagation between lhs and output along the batch dimension such that we can support data parallelism.

Add a test case, which fails before this cl and succeeds with this cl.

PiperOrigin-RevId: 614738671
@copybara-service copybara-service bot merged commit b219bdc into master Mar 11, 2024
@copybara-service copybara-service bot deleted the test_614049496 branch March 11, 2024 18:52
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.

1 participant