Unverified Commit f3c7941e authored by Kai Song's avatar Kai Song Committed by GitHub
Browse files

[Bugfix]Fix EP precision for Qwen3.5, Qwen3-Next (#39181)


Signed-off-by: default avatarSong Kai <songkai05@baidu.com>
parent 3352bf8b
...@@ -80,6 +80,7 @@ class Qwen2MoeMLP(nn.Module): ...@@ -80,6 +80,7 @@ class Qwen2MoeMLP(nn.Module):
quant_config: QuantizationConfig | None = None, quant_config: QuantizationConfig | None = None,
reduce_results: bool = True, reduce_results: bool = True,
expert_gate: torch.nn.Linear | None = None, expert_gate: torch.nn.Linear | None = None,
is_sequence_parallel: bool = False,
prefix: str = "", prefix: str = "",
) -> None: ) -> None:
super().__init__() super().__init__()
...@@ -88,6 +89,7 @@ class Qwen2MoeMLP(nn.Module): ...@@ -88,6 +89,7 @@ class Qwen2MoeMLP(nn.Module):
[intermediate_size] * 2, [intermediate_size] * 2,
bias=False, bias=False,
quant_config=quant_config, quant_config=quant_config,
disable_tp=is_sequence_parallel,
prefix=f"{prefix}.gate_up_proj", prefix=f"{prefix}.gate_up_proj",
) )
self.down_proj = RowParallelLinear( self.down_proj = RowParallelLinear(
...@@ -96,6 +98,7 @@ class Qwen2MoeMLP(nn.Module): ...@@ -96,6 +98,7 @@ class Qwen2MoeMLP(nn.Module):
bias=False, bias=False,
quant_config=quant_config, quant_config=quant_config,
reduce_results=reduce_results, reduce_results=reduce_results,
disable_tp=is_sequence_parallel,
prefix=f"{prefix}.down_proj", prefix=f"{prefix}.down_proj",
) )
if hidden_act != "silu": if hidden_act != "silu":
......
...@@ -140,6 +140,7 @@ class Qwen3NextSparseMoeBlock(nn.Module): ...@@ -140,6 +140,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
quant_config=quant_config, quant_config=quant_config,
reduce_results=False, reduce_results=False,
expert_gate=self.shared_expert_gate, expert_gate=self.shared_expert_gate,
is_sequence_parallel=self.is_sequence_parallel,
prefix=f"{prefix}.shared_expert", prefix=f"{prefix}.shared_expert",
) )
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment