You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Is your feature request related to a problem? Please describe.
When distilling between teacher and student models from the same family but with slightly different vocabulary sizes, we currently don't see a supported way to handle the mismatch in modelopt.torch.distill.
In our case, Qwen2.5-Coder-7B-Instruct has vocab_size=152,064, while Qwen2.5-Coder-0.5B-Instruct has vocab_size=151,936. LogitsDistillationLoss.forward applies F.log_softmax/F.softmax directly to the teacher and student logits and then computes F.kl_div, which requires the last dimensions to match.
KDLossConfig also doesn't appear to provide an option for configuring vocabulary alignment. As a workaround, we implemented a custom loss wrapper that truncates both logits tensors to min(teacher_vocab, student_vocab) before passing them to the criterion. This works for our case, but we weren't able to find documented guidance on whether this is the intended approach or whether it is numerically correct when used together with StaticLossBalancer.
Describe the solution you'd like
It would be helpful to have a native and documented way to handle vocabulary-size mismatches between teacher and student models.
For example, this could be:
Automatic vocabulary alignment/truncation inside LogitsDistillationLoss, or
A separate loss/configuration option specifically for vocabulary alignment, or
At minimum, a documented parameter or recommended pattern for supplying a pre-aligned vocabulary range.
This would make it easier for users to handle models from the same family without having to implement custom loss logic based on the current source code.
Describe alternatives you've considered
Custom LogitsDistillationLoss wrapper: We currently truncate both logits tensors to min(teacher_vocab, student_vocab) before calling the original criterion. This works in our setup, but we are unsure whether this is the recommended approach, particularly with StaticLossBalancer (see StaticLossBalancer warns "weights do not sum to 1.0" with a config that does sum to 1.0 #2488).
Custom PyTorch KD loop: We also considered bypassing ModelOpt's distillation wrapper and implementing the training loop ourselves. This ultimately worked for our project, but it means giving up ModelOpt's teacher/student orchestration, mtd.convert()/mtd.export() lifecycle, and the downstream quantization/export integration.
Additional context
This came up while distilling models from the same Qwen2.5-Coder family rather than from different architectures. The tokenizer vocabulary is effectively the same, and the difference appears to come from padding the embedding matrix for GPU efficiency: the 7B model uses 152,064 entries, while the 0.5B model uses 151,936.
Because vocabulary-size differences like this can occur between different model sizes within the same family, having a supported way to handle them could be useful for other distillation use cases as well.
As a separate note, we also encountered a crash when calling restore() during an earlier iteration of our script. Unfortunately, we can no longer reproduce it because that version of the script was lost, so we don't have a minimal reproduction to provide. I'm mentioning it only in case it is already a known issue.
Thank you for considering this. Any guidance on the recommended approach for handling vocabulary mismatches would also be greatly appreciated.
Is your feature request related to a problem? Please describe.
When distilling between teacher and student models from the same family but with slightly different vocabulary sizes, we currently don't see a supported way to handle the mismatch in
modelopt.torch.distill.In our case,
Qwen2.5-Coder-7B-Instructhasvocab_size=152,064, whileQwen2.5-Coder-0.5B-Instructhasvocab_size=151,936.LogitsDistillationLoss.forwardappliesF.log_softmax/F.softmaxdirectly to the teacher and student logits and then computesF.kl_div, which requires the last dimensions to match.KDLossConfigalso doesn't appear to provide an option for configuring vocabulary alignment. As a workaround, we implemented a custom loss wrapper that truncates both logits tensors tomin(teacher_vocab, student_vocab)before passing them to the criterion. This works for our case, but we weren't able to find documented guidance on whether this is the intended approach or whether it is numerically correct when used together withStaticLossBalancer.Describe the solution you'd like
It would be helpful to have a native and documented way to handle vocabulary-size mismatches between teacher and student models.
For example, this could be:
LogitsDistillationLoss, orThis would make it easier for users to handle models from the same family without having to implement custom loss logic based on the current source code.
Describe alternatives you've considered
LogitsDistillationLosswrapper: We currently truncate both logits tensors tomin(teacher_vocab, student_vocab)before calling the original criterion. This works in our setup, but we are unsure whether this is the recommended approach, particularly withStaticLossBalancer(see StaticLossBalancer warns "weights do not sum to 1.0" with a config that does sum to 1.0 #2488).mtd.convert()/mtd.export()lifecycle, and the downstream quantization/export integration.Additional context
This came up while distilling models from the same Qwen2.5-Coder family rather than from different architectures. The tokenizer vocabulary is effectively the same, and the difference appears to come from padding the embedding matrix for GPU efficiency: the 7B model uses 152,064 entries, while the 0.5B model uses 151,936.
Because vocabulary-size differences like this can occur between different model sizes within the same family, having a supported way to handle them could be useful for other distillation use cases as well.
As a separate note, we also encountered a crash when calling
restore()during an earlier iteration of our script. Unfortunately, we can no longer reproduce it because that version of the script was lost, so we don't have a minimal reproduction to provide. I'm mentioning it only in case it is already a known issue.Thank you for considering this. Any guidance on the recommended approach for handling vocabulary mismatches would also be greatly appreciated.