Use attn_tp_group for all reduce in token embedding (#17403)

2026-01-20 23:56:21 +08:00
parent ce2d686e94
commit d97066d209
4 changed files with 17 additions and 5 deletions
@@ -2474,7 +2474,7 @@ class DeepseekV2Model(nn.Module):
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
-                enable_tp=not is_dp_attention_enabled(),
+                use_attn_tp_group=is_dp_attention_enabled(),
            )
        else:
            self.embed_tokens = PPMissingLayer()