[Bugfix] Fix spec decode memory regression after #28549 (#28819)

Signed-off-by: zhewenli <zhewenli@meta.com>
This commit is contained in:
Zhewen Li
2025-11-20 03:05:50 -08:00
committed by GitHub
parent 371b1d4c61
commit 93c8672ceb
4 changed files with 5 additions and 19 deletions

View File

@@ -9,7 +9,6 @@ from transformers import LlamaConfig
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -155,10 +154,6 @@ class LlamaModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
# if PP disabled then draft will share embed with target
if get_pp_group().world_size == 1 and "embed_tokens." in name:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)