Skip to content

Fix bugs in initial_load_in_hf when enable_weight_tying=true in Qwen3 #1964

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

Closed
wants to merge 6 commits into from
Closed

Fix bugs in initial_load_in_hf when enable_weight_tying=true in Qwen3 #1964

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d57d3a0
Fix bugs in initial_load_in_hf when enable_weight_tying=true in Qwen3
Achazwl Oct 29, 2025
a005f0f
Update torchtitan/models/qwen3/model/state_dict_adapter.py
Achazwl Oct 30, 2025
978746f
Update torchtitan/models/qwen3/model/state_dict_adapter.py
Achazwl Oct 30, 2025
63d5b5f
Remove the shadow copy of hf_state_dict
Achazwl Oct 31, 2025
8a23f6e
fix lint
Achazwl Nov 4, 2025
bb5db8e
Merge branch 'main' of https://github.com/pytorch/torchtitan into pat…
Achazwl Nov 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
fix lint
  • Loading branch information
@Achazwl
Achazwl committed Nov 4, 2025
commit 8a23f6ee2956bc8f453b1522e248d8d1b1841a52
9 changes: 7 additions & 2 deletions torchtitan/models/qwen3/model/state_dict_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,14 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]:
state_dict = {}
expert_weights_by_layer = {} # {layer: {abstract_key: {expert_id: tensor}}}

if self.model_args.enable_weight_tying and "lm_head.weight" not in hf_state_dict:
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["lm_head.weight"] = hf_state_dict["model.embed_tokens.weight"]
hf_state_dict["lm_head.weight"] = hf_state_dict[
"model.embed_tokens.weight"
]

for key, value in hf_state_dict.items():
if "mlp.experts" in key:
Expand Down
Loading