New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix incorrect training of first and last Medusa heads #325

Open

chiwanpark wants to merge 1 commit into linkedin:main from chiwanpark:fix-medusa-loss

Contributor

chiwanpark commented Oct 26, 2024

Summary

Currently, there are two errors on Medusa training examples:

When we use Liger Kernel, the first head (model.medusa_head[0]) is not trained.
When we don't use Liger Kernel, the logits of the last head (medusa_logits[-1]) is ignored.

This PR fixes these errors.

Testing Done

Hardware Type: A100 80GB 8 GPUs
run make test to ensure correctness
run make checkstyle to ensure code style
run make test-convergence to ensure convergence


          Fix incorrect training of first and last Medusa heads

b1d12d6

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet