[BUG] The problems with bucket and shared_embedding.

[Baibaifan]

Describe the bug

As shown in the figure above, shared_embedding and other parameters are distinguished when building the bucket. When the data_end_index of the parameter before shared_embedding is not divisible by self.data_parallel_world_size, there will be a problem here.

To Reproduce

Expected behavior

Stack trace/logs

Environment (please complete the following information):

• Megatron-LM commit ID :c3677e09aa4e2eec37048307bd795928b8f8324a
• PyTorch version: 2.2.0a0+81ea7a4
• CUDA version: -
• NCCL version: -

Proposed fix

Additional context

[wangxicoding]

I fixed it, but the megatron has not responded👀
#762

[Baibaifan]

I fixed it, but the megatron has not responded👀 #762

It’s so awesome, I must give u a Turing Award.