From 54e964dc0ac2f717e37e33037a956c2f6a8f738f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zhengyuan=20Su=20=28=E8=8B=8F=E6=94=BF=E6=B8=8A=29?= Date: Tue, 7 Apr 2026 07:28:13 +0800 Subject: [PATCH] [Bugfix] Fix load_weights fallback for non-fused stacked_params_mapping entries (#2523) --- .../diffusion/models/bagel/bagel_transformer.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index bbcd09dd5..a14e875c0 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -944,26 +944,27 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): loaded_params: set[str] = set() for name, loaded_weight in weights: - original_name = name + loaded = False for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue - name = name.replace(weight_name, param_name) - param = params_dict.get(name) + stacked_name = name.replace(weight_name, param_name) + param = params_dict.get(stacked_name) if param is None: - # Fused param doesn't exist (e.g. gate_up_proj on DiT); - # restore original name and fall through to non-stacked path. - name = original_name break weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight, shard_id) + name = stacked_name + loaded = True break - else: + + if not loaded: param = params_dict.get(name) if param is None: continue weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) + loaded_params.add(name) return loaded_params