Skip to content

Commit 0246048

Browse files
chiajunglienchiajunglien
committed
fix(lora): refactor lora device retrieval
1 parent 636ca8a commit 0246048

5 files changed

Lines changed: 97 additions & 104 deletions

File tree

‎vllm/lora/layers/logits_processor.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,6 @@ def _get_logits(
158158

159159
if logits is None:
160160
return None
161-
162-
if current_platform.is_tpu():
163-
return logits[:, : self.base_layer.vocab_size]
164161

165162
if self.sharded_to_full_mapping_gpu is not None:
166163
# Reindex full logits tensor to ensure 1:1 mapping between

‎vllm/lora/layers/utils.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ def __post_init__(self):
3232
def _get_lora_device(base_layer: nn.Module) -> torch.device:
3333
# code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34
3434
"""Returns the device for where to place the LoRA tensors."""
35-
3635
def get_dev(obj):
3736
if obj is None:
3837
return None

‎vllm/lora/layers/vocal_parallel_embedding.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,6 @@ def set_lora(
9494
)
9595

9696
def forward(self, x: torch.Tensor) -> torch.Tensor:
97-
full_output = self.base_layer.forward(x)
98-
if current_platform.is_tpu():
99-
return full_output
10097
# NB: Don't use torch.narrow here. torch.narrow triggers some
10198
# Dynamic Shape specialization in torch.compile
10299
num_tokens = x.shape[0]
@@ -106,6 +103,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
106103
x + indices_1,
107104
self.lora_a_stacked_2d,
108105
)
106+
full_output = self.base_layer.forward(x)
109107

110108
full_output_org = full_output
111109
if full_output.ndim == 3:

‎vllm/lora/model_manager.py‎

Lines changed: 93 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from typing import TypeVar
77

88
import torch
9+
import torchax
910
from torch import nn
1011

1112
from vllm.config import VllmConfig
@@ -481,101 +482,102 @@ def create_dummy_lora(
481482
) -> LoRAModel:
482483
"""Create zero-initialized LoRAModel for warmup."""
483484
model = LoRAModel(lora_id, rank, {})
484-
for module_name, module in self.model.named_modules():
485-
if (
486-
not self._match_target_modules(module_name)
487-
or not isinstance(module, BaseLayerWithLoRA)
488-
or self._get_punica_wrapper(module_name) is None
489-
):
490-
continue
491-
parts = module_name.split(".")
492-
if module_name not in self.packed_modules:
493-
assert embedding_modules is not None
494-
if parts[-1] in embedding_modules:
495-
# Special-case lm_head: wrapped by LogitsProcessorWithLoRA.
496-
# LoRA input dim is hidden_size, output dim is vocab size.
497-
# LogitsProcessorWithLoRA handles extra vocab size directly.
498-
if parts[-1] == "lm_head":
499-
input_dim = module.lora_a_stacked[0].shape[-1]
500-
output_dim = module.lora_b_stacked[0].shape[-2]
501-
else:
502-
input_dim = (
503-
module.base_layer.org_vocab_size
504-
if hasattr(module.base_layer, "org_vocab_size")
505-
else module.base_layer.weight.shape[1]
485+
with torchax.default_env():
486+
for module_name, module in self.model.named_modules():
487+
if (
488+
not self._match_target_modules(module_name)
489+
or not isinstance(module, BaseLayerWithLoRA)
490+
or self._get_punica_wrapper(module_name) is None
491+
):
492+
continue
493+
parts = module_name.split(".")
494+
if module_name not in self.packed_modules:
495+
assert embedding_modules is not None
496+
if parts[-1] in embedding_modules:
497+
# Special-case lm_head: wrapped by LogitsProcessorWithLoRA.
498+
# LoRA input dim is hidden_size, output dim is vocab size.
499+
# LogitsProcessorWithLoRA handles extra vocab size directly.
500+
if parts[-1] == "lm_head":
501+
input_dim = module.lora_a_stacked[0].shape[-1]
502+
output_dim = module.lora_b_stacked[0].shape[-2]
503+
else:
504+
input_dim = (
505+
module.base_layer.org_vocab_size
506+
if hasattr(module.base_layer, "org_vocab_size")
507+
else module.base_layer.weight.shape[1]
508+
)
509+
output_dim = (
510+
module.base_layer.embedding_dim
511+
if hasattr(module.base_layer, "embedding_dim")
512+
else module.base_layer.weight.shape[0]
513+
)
514+
lora = LoRALayerWeights.create_dummy_lora_weights(
515+
module_name,
516+
input_dim,
517+
output_dim,
518+
rank,
519+
module.lora_a_stacked[0].dtype,
520+
"cpu",
506521
)
507-
output_dim = (
508-
module.base_layer.embedding_dim
509-
if hasattr(module.base_layer, "embedding_dim")
510-
else module.base_layer.weight.shape[0]
522+
model.loras[module_name] = lora
523+
elif module.__class__.__name__ == "FusedMoE3DWithLoRA":
524+
# Case for 3D moe model
525+
# w2
526+
lora = LoRALayerWeights.create_dummy_lora_weights(
527+
module_name,
528+
module.w2_input_size,
529+
module.w2_output_size,
530+
rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts
531+
module.w2_lora_a_stacked[0].dtype,
532+
"cpu",
511533
)
512-
lora = LoRALayerWeights.create_dummy_lora_weights(
513-
module_name,
514-
input_dim,
515-
output_dim,
516-
rank,
517-
module.lora_a_stacked[0].dtype,
518-
"cpu",
519-
)
520-
model.loras[module_name] = lora
521-
elif module.__class__.__name__ == "FusedMoE3DWithLoRA":
522-
# Case for 3D moe model
523-
# w2
524-
lora = LoRALayerWeights.create_dummy_lora_weights(
525-
module_name,
526-
module.w2_input_size,
527-
module.w2_output_size,
528-
rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts
529-
module.w2_lora_a_stacked[0].dtype,
530-
"cpu",
531-
)
532-
model.loras[module_name] = lora
533-
# w13
534-
lora = LoRALayerWeights.create_dummy_lora_weights(
535-
module_name,
536-
module.w13_input_size,
537-
module.w13_output_size,
538-
rank
539-
* module.w13_lora_a_stacked[0].shape[1], # rank*num_experts
540-
module.w13_lora_a_stacked[0].dtype,
541-
"cpu",
542-
)
543-
model.loras[module_name + ".base_layer"] = lora
534+
model.loras[module_name] = lora
535+
# w13
536+
lora = LoRALayerWeights.create_dummy_lora_weights(
537+
module_name,
538+
module.w13_input_size,
539+
module.w13_output_size,
540+
rank
541+
* module.w13_lora_a_stacked[0].shape[1], # rank*num_experts
542+
module.w13_lora_a_stacked[0].dtype,
543+
"cpu",
544+
)
545+
model.loras[module_name + ".base_layer"] = lora
546+
else:
547+
lora = LoRALayerWeights.create_dummy_lora_weights(
548+
module_name,
549+
module.lora_a_stacked[0].shape[-1],
550+
module.lora_b_stacked[0].shape[-2],
551+
rank,
552+
module.lora_a_stacked[0].dtype,
553+
"cpu",
554+
)
555+
model.loras[module_name] = lora
544556
else:
545-
lora = LoRALayerWeights.create_dummy_lora_weights(
546-
module_name,
547-
module.lora_a_stacked[0].shape[-1],
548-
module.lora_b_stacked[0].shape[-2],
549-
rank,
550-
module.lora_a_stacked[0].dtype,
551-
"cpu",
552-
)
557+
parts = module_name.split(".")
558+
replacements = self.packed_modules_mapping[parts[-1]]
559+
subloras: list[LoRALayerWeights | None] = []
560+
for i, r in enumerate(replacements):
561+
lora = LoRALayerWeights.create_dummy_lora_weights(
562+
module_name + "." + r,
563+
module.lora_a_stacked[i].shape[-1],
564+
module.lora_b_stacked[i].shape[-2],
565+
rank,
566+
module.lora_a_stacked[i].dtype,
567+
"cpu",
568+
)
569+
subloras.append(lora)
570+
if module.__class__.__name__ == "FusedMoEWithLoRA":
571+
# For non-gated MoE, pad subloras to 3 elements per expert
572+
# to match pack_moe expectations (w1, w2, None for w3)
573+
if self._is_non_gated_moe and len(subloras) > 0:
574+
subloras = self._pad_lora_pairs_to_triplets(subloras)
575+
lora = PackedLoRALayerWeights.pack_moe(
576+
subloras, module_name, is_non_gated_moe=self._is_non_gated_moe
577+
)
578+
else:
579+
lora = PackedLoRALayerWeights.pack(subloras)
553580
model.loras[module_name] = lora
554-
else:
555-
parts = module_name.split(".")
556-
replacements = self.packed_modules_mapping[parts[-1]]
557-
subloras: list[LoRALayerWeights | None] = []
558-
for i, r in enumerate(replacements):
559-
lora = LoRALayerWeights.create_dummy_lora_weights(
560-
module_name + "." + r,
561-
module.lora_a_stacked[i].shape[-1],
562-
module.lora_b_stacked[i].shape[-2],
563-
rank,
564-
module.lora_a_stacked[i].dtype,
565-
"cpu",
566-
)
567-
subloras.append(lora)
568-
if module.__class__.__name__ == "FusedMoEWithLoRA":
569-
# For non-gated MoE, pad subloras to 3 elements per expert
570-
# to match pack_moe expectations (w1, w2, None for w3)
571-
if self._is_non_gated_moe and len(subloras) > 0:
572-
subloras = self._pad_lora_pairs_to_triplets(subloras)
573-
lora = PackedLoRALayerWeights.pack_moe(
574-
subloras, module_name, is_non_gated_moe=self._is_non_gated_moe
575-
)
576-
else:
577-
lora = PackedLoRALayerWeights.pack(subloras)
578-
model.loras[module_name] = lora
579581
return model
580582

581583
def _match_target_modules(self, module_name: str) -> bool:

‎vllm/model_executor/model_loader/default_loader.py‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -308,12 +308,9 @@ def _init_ep_weight_filter(self, model_config: ModelConfig) -> None:
308308
expert weights. By computing the set upfront we can skip non-local
309309
expert tensors *before* reading them from disk.
310310
"""
311-
try:
312-
from vllm.config import get_current_vllm_config
313-
vllm_config = get_current_vllm_config()
314-
except AssertionError:
315-
logger.warning("vLLM config not set during weight filtering, skipping EP filter.")
316-
return
311+
from vllm.config import get_current_vllm_config
312+
313+
vllm_config = get_current_vllm_config()
317314
parallel_config = vllm_config.parallel_config
318315

319316
if not (

0 commit comments

Comments
 (0)