Wrong output fp16/bf16 dtype in ParallelEmbedding when sharding accross vocab

[dacorvo]

In the ParallelEmbedding layer, when sharding accross vocab, the output is masked at the very end of the operation.

It seems that the masking is done by multiplying by an hard-coded float mask, which leads to the actual float16/bfloat16 to be upcast to float32.

A correct implementation would be to multiply by a mask of the same type as the intended output.

[fayyadd]

thanks for reporting! Our team will take a look.

[fayyadd]

Thanks for reaching out. Our team has a fix for this and it will be available in the upcoming releases. We will update this issue once the fix is released.