mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> This PR fixes #309. The original implementation injects Medusa heads to FSDP-unwrapped model variables, which are not accessed from the state_dict (FSDP-wrapped). To fix this bug, I've changed the method to inject Medusa heads with model loader, and the method to access the state_dict through APIs from accelerate. Minor updates: - Update dependencies to the latest. - Fix scripts to run Medusa example. <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. --->  <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: A100-80G-NVLINK 8 cards - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Co-authored-by: Shao Tang <[email protected]>