[Bugfix] Shard embedding hooks

[comaniac]

Description

When sharding word embedding, we use forward pre-hook to mask inputs for different devices in a TP group. Accordingly, we have to mask the invalid output using the same mask in forward post-hook. Since the input mask generated in pre-hook cannot be saved, we currently re-generate the mask in post-hook. However, at this moment the input is already masked, so we cannot re-generate a mask using the same logic.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @chhzh123 @zarzen

[chhzh123]

Good catch. This bug is pretty subtle. Thanks for fixing.

[comaniac]

Thanks @chhzh123