[Qwen3.5] Fix broken pipeline parallelism layer splitting (#21070)

Co-authored-by: Alison Shao <alison.shao@Mac.attlocal.net>
This commit is contained in:
Alison Shao
2026-03-21 01:02:51 -07:00
committed by GitHub
parent 9ca68a5904
commit 852e112ebf
2 changed files with 15 additions and 20 deletions

View File

@@ -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

View File

@@ -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%}"
),
)