Add sharding of conv2d, group norm and layer norm - Githubissues

nod-ai / sharktank

SHARK Inference Modeling and Serving

Apache License 2.0

7 stars 9 forks source link

Add sharding of conv2d, group norm and layer norm #50

Closed sogartar closed 1 month ago

sogartar commented 1 month ago

Sharding of output channel dimension of conv2d with a single group.
Sharding by groups in group normalization.
Sharding of non-normalization dimension in layer norm.

sogartar commented 1 month ago

Rebased before merging.