Inconsistent loss when resume training with vocab size that is not divisible by world size.

[weixuansun]

Bug description

When I use an tokenizer that has vocabulary size that is not divisible by parallel (or world) size, the training loss will become inconsistent after resuming.

Versions

can be reproduced using torch2.6

Reproduce:
Use any tokenizer that has a vocabulary size that is not divisible by parallel (or world) size.
train from step 0 to step 20:

load step 10 checkpoint and resume training: As shown, step 11 and following steps have inconsistent loss.

[tianyu-l]

could you try the latest pytorch nightly? might have been fixed by pytorch/pytorch#150490

[unlimblue]

could you try the latest pytorch nightly? might have been fixed by pytorch/pytorch#150490

https://github.com/pytorch/pytorch/tree/c729f7dbee3be20f04a6aeac41ae0cd5be23403d

Using the latest main branch source code from April 22nd to compile PyTorch still results in the same issue.

[fegin]

This looks like a checkpointing problem. Can you confirm 1) does this only happen after checkpoint resume? 2) does this only happen when vocab size is not divisible by world size? From your description, the answers to both questions seem to be yes, but I do want to confirm. Thanks!

[weixuansun]

This looks like a checkpointing problem. Can you confirm 1) does this only happen after checkpoint resume? 2) does this only happen when vocab size is not divisible by world size? From your description, the answers to both questions seem to be yes, but I do want to confirm. Thanks!

Yes, answers to both questions are Yes.
If I pad the vocab size to be divisible, the loss would be consistent.