Fix qwen3 tbo/dp-lm-head (#6652)

2025-05-27 00:38:27 -07:00
parent ce9d690ef4
commit b18416fbf8
3 changed files with 3 additions and 1 deletions
@@ -688,6 +688,7 @@ class Qwen3MoeForCausalLM(nn.Module):
            config.hidden_size,
            quant_config=quant_config,
            prefix=add_prefix("lm_head", prefix),
+            use_attn_tp_group=global_server_args_dict["enable_dp_lm_head"],
        )
        self.logits_processor = LogitsProcessor(config)