diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 0cf1b4bf9..107b73378 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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 diff --git a/test/registered/distributed/test_pp_single_node.py b/test/registered/distributed/test_pp_single_node.py index 2468b82b7..6b0f4557e 100644 --- a/test/registered/distributed/test_pp_single_node.py +++ b/test/registered/distributed/test_pp_single_node.py @@ -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%}" ), )