diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 30b02ec83..0760c4645 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -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 diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index a3d5fc4a4..f0ee994a1 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -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: