[Qwen3.5] Fix broken pipeline parallelism layer splitting (#21070)
Co-authored-by: Alison Shao <alison.shao@Mac.attlocal.net>
This commit is contained in:
@@ -29,7 +29,7 @@ from sglang.srt.configs.qwen3_5 import (
|
||||
)
|
||||
|
||||
# Distributed
|
||||
from sglang.srt.distributed import get_pp_group, get_pp_indices
|
||||
from sglang.srt.distributed import get_pp_group
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
|
||||
|
||||
@@ -721,25 +721,14 @@ class Qwen3_5ForCausalLM(nn.Module):
|
||||
is_nextn=is_nextn,
|
||||
)
|
||||
|
||||
self.layers = make_layers(
|
||||
self.layers, self._start_layer, self._end_layer = make_layers(
|
||||
config.num_hidden_layers,
|
||||
get_layer,
|
||||
pp_rank=self.pp_group.rank_in_group,
|
||||
pp_size=self.pp_group.world_size,
|
||||
prefix=f"{prefix}.layers",
|
||||
)
|
||||
|
||||
pp_rank = self.pp_group.rank_in_group
|
||||
pp_size = self.pp_group.world_size
|
||||
num_layers = config.num_hidden_layers
|
||||
self._start_layer, self._end_layer = (
|
||||
get_pp_indices(
|
||||
num_layers,
|
||||
pp_rank,
|
||||
pp_size,
|
||||
)
|
||||
if pp_rank is not None and pp_size is not None
|
||||
else (0, num_layers)
|
||||
)
|
||||
|
||||
# Final normalization
|
||||
if self.pp_group.is_last_rank:
|
||||
self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
@@ -945,6 +934,8 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
|
||||
shard_id: str,
|
||||
num_experts: int,
|
||||
):
|
||||
if name not in params_dict:
|
||||
return False
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
# let ep moe layer to gracefully handle expert_ids that do not belong to local moe rank
|
||||
@@ -1270,6 +1261,8 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
|
||||
shard_id: str,
|
||||
num_experts: int,
|
||||
):
|
||||
if name not in params_dict:
|
||||
return False
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
# let ep moe layer to gracefully handle expert_ids that do not belong to local moe rank
|
||||
|
||||
@@ -358,12 +358,14 @@ class TestQwen35PPAccuracy(unittest.TestCase):
|
||||
"Qwen/Qwen3.5-35B-A3B" # replace with your Qwen Model if needed
|
||||
)
|
||||
|
||||
def run_gsm8k_test(self, pp_size):
|
||||
def run_gsm8k_test(self, tp_size, pp_size):
|
||||
process = popen_launch_server(
|
||||
self.model_name,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
tp_size,
|
||||
"--pp-size",
|
||||
pp_size,
|
||||
"--chunked-prefill-size",
|
||||
@@ -388,17 +390,17 @@ class TestQwen35PPAccuracy(unittest.TestCase):
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def test_pp_consistency(self):
|
||||
baseline = self.run_gsm8k_test(pp_size=1)
|
||||
pp_metrics = self.run_gsm8k_test(pp_size=2)
|
||||
baseline = self.run_gsm8k_test(tp_size=2, pp_size=1)
|
||||
pp_metrics = self.run_gsm8k_test(tp_size=1, pp_size=2)
|
||||
|
||||
print(f"[Qwen35 PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
|
||||
|
||||
self.assertGreaterEqual(baseline["accuracy"], 0.83)
|
||||
self.assertGreaterEqual(
|
||||
pp_metrics["accuracy"],
|
||||
baseline["accuracy"] - 0.02,
|
||||
baseline["accuracy"] - 0.05,
|
||||
msg=(
|
||||
f"PP accuracy dropped more than 2% compared to baseline. "
|
||||
f"PP accuracy dropped more than 5% compared to baseline. "
|
||||
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user