A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Disable `torch.compile` for dynamic rope models in Transformers backend (#23738)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 27, 2025 at 20:03 UTC
0585a9e73c072a8cbb1a64bea3c26dd0d2dde402
1 file changed
+22
-3
vllm/model_executor/models/transformers.py
+22
-3
@@ -88,6 +88,23 @@ def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):
88
logger.debug("%s: %s -> %s", name, old_module, new_module)
89
90
91
+def can_enable_torch_compile(vllm_config: VllmConfig) -> bool:
92
+ """
93
+ Callable to be passed to `@support_torch_compile`'s `enable_if` argument.
94
+
95
+ Defaults to `True` but is disabled in the following situations:
96
+
97
+ - The model uses dynamic rope scaling.
98
+ """
99
+ enable = True
100
+ text_config = vllm_config.model_config.hf_config.get_text_config()
101
+ # Dynamic rope scaling is not compatible with torch.compile
102
+ rope_scaling: dict = getattr(text_config, "rope_scaling", None) or {}
103
+ if rope_scaling.get("rope_type") == "dynamic":
104
+ enable = False
105
+ return enable
106
+
107
+
108
def replace_linear_class(
109
linear: nn.Linear, style: Literal["colwise", "rowwise"],
110
quant_config: QuantizationConfig
@@ -641,7 +658,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
658
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
659
660
644
-@support_torch_compile
661
+@support_torch_compile(enable_if=can_enable_torch_compile)
662
class TransformersModel(TransformersBase):
663
hf_to_vllm_mapper = WeightsMapper(
664
orig_to_new_prefix={
@@ -653,7 +670,7 @@ class TransformersModel(TransformersBase):
670
})
671
672
656
-@support_torch_compile
673
+@support_torch_compile(enable_if=can_enable_torch_compile)
674
class TransformersForCausalLM(TransformersBase):
675
676
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -709,12 +726,14 @@ def flatten_and_concat(x: list[torch.Tensor]) -> torch.Tensor:
726
info=MultiModalProcessingInfo,
727
dummy_inputs=MultiModalDummyInputsBuilder)
728
@support_torch_compile(
729
+ # set `positions` to last dim to support Qwen-mrope
730
dynamic_arg_dims={
731
"input_ids": 0,
732
"positions": -1,
733
"intermediate_tensors": 0,
734
"inputs_embeds": 0,
717
- }) # set `positions` to last dim to support Qwen-mrope
735
+ },
736
+ enable_if=can_enable_torch_compile)
737
class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):
738
# Backwards compatibility for prev released models. State dicts back then
739
# had different formats and cannot be loaded with `AutoModel` mapping as is