[bug fix][pp] fix qwen3 model load (#15223)
Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
@@ -392,13 +392,13 @@ class Qwen3ForCausalLM(nn.Module):
|
||||
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
|
||||
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=(config.vocab_size, config.hidden_size),
|
||||
size=self.lm_head.weight.shape,
|
||||
dtype=next(self.model.parameters()).dtype,
|
||||
src=self.pp_group.first_rank,
|
||||
src=0,
|
||||
)
|
||||
self.lm_head.weight.copy_(emb_token_weight)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user