chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PP] fix wrong weight logic for tie_word_embeddings model (#15890)

Browse Source 
Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
Xuchun Shang
2026-01-27 17:41:17 +08:00
committed by GitHub
parent 7106f6c8e1
commit dba264ac73
2 changed files with 19 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/sglang/srt/models/qwen2.py
View File

@@ -457,20 +457,6 @@ class Qwen2ForCausalLM(nn.Module):
# ranks other than the last rank will have a placeholder layer
self.lm_head = PPMissingLayer()
# perform weight tying for PP
if self.pp_group.world_size > 1 and config.tie_word_embeddings:
if self.pp_group.is_first_rank:
self.pp_group.send(
self.model.embed_tokens.weight, dst=self.pp_group.last_rank
)
elif self.pp_group.is_last_rank:
emb_token_weight = self.pp_group.recv(
size=(config.vocab_size, config.hidden_size),
dtype=next(self.model.parameters()).dtype,
src=self.pp_group.first_rank,
)
self.lm_head.weight.copy_(emb_token_weight)
self.logits_processor = LogitsProcessor(config)
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
# For EAGLE3 support
@@ -589,22 +575,21 @@ class Qwen2ForCausalLM(nn.Module):
):
continue
if name == "model.embed_tokens.weight":
if self.pp_group.is_last_rank and self.config.tie_word_embeddings:
if "lm_head.weight" in params_dict:
param = params_dict["lm_head.weight"]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
if "rotary_emb.inv_freq" in name or "projector" in name:
continue
if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
# Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
continue
if self.config.tie_word_embeddings and "lm_head.weight" in name:
if self.pp_group.world_size > 1 and self.pp_group.is_last_rank:
# Handle pp weight tying here
# find the embed_tokens.weight in the weights
embed_token_weights = next(
filter(lambda x: x[0] == "model.embed_tokens.weight", weights)
)[1]
loaded_weight = embed_token_weights
else:
continue
if name.startswith("model.vision_tower") and name not in params_dict:
continue

34
python/sglang/srt/models/qwen3.py
View File

@@ -379,20 +379,6 @@ class Qwen3ForCausalLM(nn.Module):
# ranks other than the last rank will have a placeholder layer
self.lm_head = PPMissingLayer()
# perform weight tying for PP
if self.pp_group.world_size > 1 and config.tie_word_embeddings:
if self.pp_group.is_first_rank:
self.pp_group.send(
self.model.embed_tokens.weight, dst=self.pp_group.world_size - 1
)
elif self.pp_group.is_last_rank:
emb_token_weight = self.pp_group.recv(
size=self.lm_head.weight.shape,
dtype=next(self.model.parameters()).dtype,
src=0,
)
self.lm_head.weight.copy_(emb_token_weight)
self.logits_processor = LogitsProcessor(config)
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
@@ -501,6 +487,16 @@ class Qwen3ForCausalLM(nn.Module):
for name, loaded_weight in weights:
if "Embedding" in self.config.name_or_path:
name = add_prefix(name, "model")
if name == "model.embed_tokens.weight":
if self.pp_group.is_last_rank and self.config.tie_word_embeddings:
if "lm_head.weight" in params_dict:
param = params_dict["lm_head.weight"]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
layer_id = get_layer_id(name)
if (
layer_id is not None
@@ -518,16 +514,6 @@ class Qwen3ForCausalLM(nn.Module):
# Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
continue
if self.config.tie_word_embeddings and "lm_head.weight" in name:
if self.pp_group.world_size > 1 and self.pp_group.is_last_rank:
# Handle pp weight tying here
# find the embed_tokens.weight in the weights
embed_token_weights = next(
filter(lambda x: x[0] == "model.embed_tokens.weight", weights)
)[1]
loaded_weight = embed_token_weights
else:
continue
if name.startswith("model.vision_tower") and name not in params_dict:
continue
if "scale" in name: