From 031d0a2aad444b61c3bbe04e67dda84677b2289d Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Tue, 10 Mar 2026 18:34:56 +0100 Subject: [PATCH] [Qwen-MOE] Fix memory duplication issues in case layers weights are re-assigned during weight loading (#18255) --- python/sglang/srt/models/qwen3_moe.py | 7 +++---- python/sglang/srt/models/qwen3_omni_moe.py | 6 ++---- python/sglang/srt/models/qwen3_vl_moe.py | 7 +++---- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 845502fe2..048a000d5 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -1043,10 +1043,9 @@ class Qwen3MoeForCausalLM(nn.Module): num_experts=self.config.num_experts, ) - # Cache params_dict to avoid repeated expensive traversal of model parameters - if not hasattr(self, "_cached_params_dict"): - self._cached_params_dict = dict(self.named_parameters()) - params_dict = self._cached_params_dict + # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: layer_id = get_layer_id(name) if ( diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index e0e82d03e..3fc5d4c6f 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/python/sglang/srt/models/qwen3_omni_moe.py @@ -524,10 +524,8 @@ class Qwen3OmniMoeForConditionalGeneration(PreTrainedModel): num_experts = self.config.num_experts - # Cache params_dict to avoid repeated expensive traversal of model parameters - if not hasattr(self, "_cached_params_dict"): - self._cached_params_dict = dict(self.named_parameters()) - params_dict = self._cached_params_dict + # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. + params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: name = name.replace(r"model.language_model.", r"model.") diff --git a/python/sglang/srt/models/qwen3_vl_moe.py b/python/sglang/srt/models/qwen3_vl_moe.py index f98460665..cf1cb3879 100644 --- a/python/sglang/srt/models/qwen3_vl_moe.py +++ b/python/sglang/srt/models/qwen3_vl_moe.py @@ -226,10 +226,9 @@ class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration): num_experts = self.config.num_experts - # Cache params_dict to avoid repeated expensive traversal of model parameters - if not hasattr(self, "_cached_params_dict"): - self._cached_params_dict = dict(self.named_parameters()) - params_dict = self._cached_params_dict + # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: name = name.replace(r"model.language_model.", r"model.")