RuntimeError: table-wise sharding on a single EmbeddingBag is not supported yet

[jiannanWang]

I saw this error message when I was trying the EmbeddingBagSharder. I read the code and found table-wise sharding is not supported in ShardedEmbeddingBag. I believe this is a bug because although table-wise sharding is not supported, the planner still generates a plan to use table-wise sharding. I found the cause for this bug might be when EmbeddingBagSharder inherits from BaseEmbeddingSharder, it doesn't override the sharding_types() function. And table-wise sharding is supported in BaseEmbeddingSharder.sharding_types(). A possible fix is to override EmbeddingBagSharder.sharding_types() function and exclude table-wise sharding. This is my first time submitting an issue to torchrec. I would appreciate any help to confirm this case.

You can find the colab here for the code to reproduce this error.

[YLGH]

hey @jiannanWang thanks for reporting this issue, we'll double check on that.

Just curious, is there any reason you're not using EmbeddingBagCollectionSharder? Table wise should be supported for it. This is the sharder we support the most (we haven't been developing on EmbeddingBagSharder)

[YLGH]

@jiannanWang I think you're diagnosis of problem, and proposed fix is correct. Feel free to make a PR!

[jiannanWang]

Hi YLGH,

Thank you for your quick response! I will make a PR shortly.

For your question about why not use EmbeddingBagCollectionSharder, I'm testing the torchrec APIs so I would like to try all the APIs provided. I saw EmbeddingBagSharder is tested in torchrec existing test cases, so I assume it's still a valid API.

[YLGH]

Closing for now - agree that this is a feature gap, but not a common one. Feel free to open a PR and we'll get it landed :)