diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index db97e2958..a6471de6b 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -83,6 +83,7 @@ from sglang.srt.runtime_context import ( get_stream, ) from sglang.srt.utils import ( + LazyValue, add_prefix, is_cuda, is_non_idle_and_non_empty, @@ -965,12 +966,17 @@ class LLaDA2MoeModelLM(nn.Module): ) weight_loader(param, loaded_weight) - self.routed_experts_weights_of_layer = { - layer_id: layer.mlp.get_moe_weights() - for layer_id, layer in enumerate(self.model.layers) - if not isinstance(layer, PPMissingLayer) - and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) - } + # Lazy: get_moe_weights() snapshots x.data, and building the map here would + # pin every expert weight's pre-process_weights_after_loading storage. + if not hasattr(self, "routed_experts_weights_of_layer"): + self.routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if not isinstance(layer, PPMissingLayer) + and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) + } + ) @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py index b8216893a..7550fd88d 100644 --- a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py +++ b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py @@ -1,4 +1,3 @@ -import os import unittest from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin @@ -31,10 +30,6 @@ class TestLLaDA2Mini(GSM8KAscendMixin, CustomTestCase): "LowConfidence", # TODO: Add dLLM configurations "--no-dllm-fdfo", # FDFO (PR #27551) halves single-batch speed on NPU; use sync mode ] - env = { - **os.environ, - "SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT": "1", # Need to avoid OOM issue - } accuracy = 0.88 output_throughput = 70