Fix bugs in initial_load_in_hf when enable_weight_tying=true in Qwen3 #1964
Conversation
Add checks for weight tying in state_dict processing
Achazwl
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
|
Hi @Achazwl! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at [email protected]. Thanks! |
meta-cla
bot
added
the
CLA Signed
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
| if key not in to_hf_map: | ||
| continue | ||
| # Skip output.weight if weight tying is enabled (HF checkpoint won't have lm_head.weight) | ||
| if self.model_args.enable_weight_tying and key == "output.weight": |
There was a problem hiding this comment.
By checking 0.6B and 1.7B model weights, they do have separate weights for embed_tokens and lm_head, and I assumed these 2 weight are the same (please correct me if I am wrong), so loading the same weights twice are ok here.
1.7B: https://huggingface.co/Qwen/Qwen3-1.7B/blob/main/model.safetensors.index.json
0.6B: https://huggingface.co/Qwen/Qwen3-0.6B/blob/main/model.safetensors
I see your change makes sense. Our previous code will fail when loading the 4B model weights: 4B model doesn't have "lm_head.weight" in their checkpoint files, but our translated hf_state_dict will still have key lm_head.weight. Did you verified the updated code still on par with HF forward? cc @shuhuayu
There was a problem hiding this comment.
Thx for catching the bug when loading the qwen3 4b model. I did a forward parity check, it works well.
| # copy from embed_tokens.weight | ||
| if self.model_args.enable_weight_tying and "lm_head.weight" not in hf_state_dict: | ||
| if "model.embed_tokens.weight" in hf_state_dict: | ||
| hf_state_dict = dict(hf_state_dict) # Make a copy to avoid modifying original |
There was a problem hiding this comment.
Do you need to make a shallow copy of the dict? Can you elaborate more on "avoid modify original"?
There was a problem hiding this comment.
Without dict(hf_state_dict), the line hf_state_dict["lm_head.weight"] = ... would directly mutate the dictionary object provided by the caller function. I'm not sure if the caller expects the input dictionary to be modified, so I made a copy to avoid any potential side effects.
There was a problem hiding this comment.
If not necessary, revert this line.
There was a problem hiding this comment.
Thanks, the input hf_state_dict will not be used after calling from_hf() function:
It should be mutate the dictionary object (hf_state_dict)
There was a problem hiding this comment.
Thanks, the input
hf_state_dictwill not be used after calling from_hf() function:It should be mutate the dictionary object (
hf_state_dict)
Ok, I've removed the shallow copy.
Co-authored-by: Shuhua Yu <[email protected]>
Co-authored-by: Shuhua Yu <[email protected]>
|
Hi @Achazwl Thanks for contribution! Do you wanna fix the lint and run CI before again before we can merge it? |
3 participants
Add checks for weight tying in state_dict processing