@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 10d765482d

Fork 0

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

`FusedMoE` support for the Transformers backend (#22650)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Oct 3, 2025 at 07:12 UTC 10d765482d19abfab6c66b5f815720a66aa9de42

10 files changed +485 -91

docs/models/supported_models.md

+3 -2

@@ -17,12 +17,12 @@ These models are what we list in [supported-text-models][supported-text-models]

18 ### Transformers

20 -vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <1% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers backend".

20 +vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers backend".

22 Currently, the Transformers backend works for the following:

24 - Modalities: embedding models, language models and vision-language models*

25 -- Architectures: encoder-only, decoder-only

25 +- Architectures: encoder-only, decoder-only, mixture-of-experts

26 - Attention types: full attention and/or sliding attention

28 _*Vision-language models currently accept only image inputs. Support for video inputs will be added in a future release._

@@ -31,6 +31,7 @@ If the Transformers model implementation follows all the steps in [writing a cus

32 - All the features listed in the [compatibility matrix](../features/README.md#feature-x-feature)

33 - Any combination of the following vLLM parallelisation schemes:

34 + - Data parallel

35 - Pipeline parallel

36 - Tensor parallel

tests/models/registry.py

@@ -661,6 +661,10 @@ _TRANSFORMERS_BACKEND_MODELS = {

661 "TransformersForSequenceClassification": _HfExamplesInfo("papluca/xlm-roberta-base-language-detection", min_transformers_version="4.57.0.dev0"), # noqa: E501

662 "TransformersForCausalLM": _HfExamplesInfo("hmellor/Ilama-3.2-1B", trust_remote_code=True), # noqa: E501

663 "TransformersForMultimodalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),

664 + "TransformersMoEForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924", min_transformers_version="4.57.0.dev0"), # noqa: E501

665 + "TransformersMoEForMultimodalLM": _HfExamplesInfo("Qwen/Qwen3-VL-30B-A3B-Instruct", min_transformers_version="4.57.0.dev0"), # noqa: E501

666 + "TransformersMoEEmbeddingModel": _HfExamplesInfo("Qwen/Qwen3-30B-A3B", min_transformers_version="4.57.0.dev0"), # noqa: E501

667 + "TransformersMoEForSequenceClassification": _HfExamplesInfo("Qwen/Qwen3-30B-A3B", min_transformers_version="4.57.0.dev0"), # noqa: E501

668 }

669

670 _EXAMPLE_MODELS = {

tests/models/test_transformers.py

@@ -66,6 +66,7 @@ def check_implementation(

66 [

67 ("meta-llama/Llama-3.2-1B-Instruct", "transformers"),

68 ("hmellor/Ilama-3.2-1B", "auto"), # CUSTOM CODE

69 + ("allenai/OLMoE-1B-7B-0924", "transformers"), # MoE

70 ]) # trust_remote_code=True by default

71 def test_models(

72 hf_runner: type[HfRunner],

@@ -74,6 +75,14 @@ def test_models(

75 model: str,

76 model_impl: str,

77 ) -> None:

78 + import transformers

79 + from packaging.version import Version

80 + installed = Version(transformers.__version__)

81 + required = Version("4.57.0.dev0")

82 + if model == "allenai/OLMoE-1B-7B-0924" and installed < required:

83 + pytest.skip("MoE models with the Transformers backend require "

84 + f"transformers>={required}, but got {installed}")

85 +

86 check_implementation(hf_runner,

87 vllm_runner,

88 example_prompts,

tests/models/utils.py

+16 -7

@@ -430,17 +430,26 @@ def dummy_hf_overrides(

430

431 update_dict = {

432 "num_layers": num_layers,

433 - "num_experts": num_experts,

434 - "num_experts_per_tok": 2,

435 - "num_local_experts": num_experts,

436 - # Otherwise there will not be any expert layers

437 - "first_k_dense_replace": 0,

438 - # To avoid OOM on DeepSeek-V3

439 - "n_routed_experts": num_experts,

433 # For Gemma-3n

434 "num_kv_shared_layers": 1,

435 }

436

437 + class DummyConfig:

438 + hf_text_config = text_config

439 +

440 + # Only set MoE related config when the model has MoE layers.

441 + # Otherwise all models detected as MoE by _get_transformers_backend_cls.

442 + if ModelConfig.get_num_experts(DummyConfig) > 0:

443 + update_dict.update({

444 + "num_experts": num_experts,

445 + "num_experts_per_tok": 2,

446 + "num_local_experts": num_experts,

447 + # Otherwise there will not be any expert layers

448 + "first_k_dense_replace": 0,

449 + # To avoid OOM on DeepSeek-V3

450 + "n_routed_experts": num_experts,

451 + })

452 +

453 # Update num_hidden_layers for non-Longcat architectures

454 if model_arch != "LongcatFlashForCausalLM" \

455 and model_arch != "LongCatFlashMTPModel":

vllm/config/model.py

+23 -16

@@ -20,7 +20,7 @@ from vllm.config.multimodal import (MMCacheType, MMEncoderTPMode,

20 MultiModalConfig)

21 from vllm.config.pooler import PoolerConfig

22 from vllm.config.scheduler import RunnerType

23 -from vllm.config.utils import assert_hashable, config

23 +from vllm.config.utils import assert_hashable, config, getattr_iter

24 from vllm.logger import init_logger

25 from vllm.platforms import current_platform

26 from vllm.transformers_utils.config import (

@@ -667,6 +667,8 @@ class ModelConfig:

667 def _get_transformers_backend_cls(self) -> str:

668 """Determine which Transformers backend class will be used if

669 `model_impl` is set to `transformers` or `auto`."""

670 + prefix = "Transformers"

671 + prefix += "MoE" if self.get_num_experts() > 1 else ""

672 # Check if the architecture we're wrapping has defaults

673 runner = None

674 convert = None

@@ -685,15 +687,15 @@ class ModelConfig:

687 # Resolve Transformers backend pooling classes

688 if runner == "pooling":

689 if convert == "embed":

688 - return "TransformersEmbeddingModel"

690 + return prefix + "EmbeddingModel"

691 if convert == "classify":

690 - return "TransformersForSequenceClassification"

692 + return prefix + "ForSequenceClassification"

693 # Resolve Transformers backend generate classes

694 if self.hf_config != self.hf_text_config:

695 # If 'hf_text_config' is the same as 'hf_config'. If not, it is

696 # probably a composite config, i.e. multimodal

695 - return "TransformersForMultimodalLM"

696 - return "TransformersForCausalLM"

697 + return prefix + "ForMultimodalLM"

698 + return prefix + "ForCausalLM"

699

700 def using_transformers_backend(self) -> bool:

701 """Check if the model is using the Transformers backend class."""

@@ -1025,17 +1027,7 @@ class ModelConfig:

1027 self.enforce_eager = True

1028

1029 def _verify_with_expert_parallelism(self) -> None:

1028 - num_expert_names = [

1029 - "moe_num_experts", # Dbrx

1030 - "num_experts", # Jamba

1031 - "n_routed_experts", # DeepSeek

1032 - "num_local_experts", # Mixtral

1033 - ]

1034 - num_experts = 0

1035 - for name in num_expert_names:

1036 - num_experts = getattr(self.hf_text_config, name, 0)

1037 - if num_experts > 0:

1038 - break

1030 + num_experts = self.get_num_experts()

1031 if num_experts < 1:

1032 raise ValueError(

1033 "Number of experts in the model must be greater than 0 "

@@ -1220,6 +1212,21 @@ class ModelConfig:

1212 num_heads = getattr(self.hf_text_config, "num_attention_heads", 0)

1213 return num_heads // parallel_config.tensor_parallel_size

1214

1215 + def get_num_experts(self) -> int:

1216 + """Returns the number of experts in the model."""

1217 + num_expert_names = [

1218 + "num_experts", # Jamba

1219 + "moe_num_experts", # Dbrx

1220 + "n_routed_experts", # DeepSeek

1221 + "num_local_experts", # Mixtral

1222 + ]

1223 + num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0)

1224 + if isinstance(num_experts, list):

1225 + # Ernie VL's remote code uses list[int]...

1226 + # The values are always the same so we just take the first one.

1227 + return num_experts[0]

1228 + return num_experts

1229 +

1230 def get_layers_start_end_indices(

1231 self, parallel_config: ParallelConfig) -> tuple[int, int]:

1232 from vllm.distributed.utils import get_pp_indices

vllm/model_executor/layers/fused_moe/layer.py

+31

@@ -960,6 +960,7 @@ class FusedMoE(CustomOp):

960 is_sequence_parallel=False,

961 zero_expert_num: Optional[int] = 0,

962 zero_expert_type: Optional[str] = None,

963 + expert_mapping: Optional[list[tuple[str, str, int, str]]] = None,

964 ):

965 super().__init__()

966 if params_dtype is None:

@@ -996,6 +997,9 @@ class FusedMoE(CustomOp):

997 self.zero_expert_num = zero_expert_num

998 self.zero_expert_type = zero_expert_type

999

1000 + # Expert mapping used in self.load_weights

1001 + self.expert_mapping = expert_mapping

1002 +

1003 # Round up hidden size if needed.

1004 hidden_size = maybe_roundup_hidden_size(hidden_size, moe_in_dtype,

1005 quant_config,

@@ -1617,6 +1621,33 @@ class FusedMoE(CustomOp):

1621

1622 return False if return_success else None

1623

1624 + def load_weights(

1625 + self, weights: Iterable[tuple[str,

1626 + torch.Tensor]]) -> Iterable[str]:

1627 + if (expert_mapping := self.expert_mapping) is None:

1628 + raise ValueError("`self.expert_mapping` must be provided to "

1629 + "load weights using `self.load_weights`.")

1630 + for expert_name, loaded_weight in weights:

1631 + qual_name = f"{self.layer_name}.{expert_name}"

1632 + for param_name, weight_name, expert_id, shard_id in expert_mapping:

1633 + if weight_name not in qual_name:

1634 + continue

1635 + weight_name = qual_name.replace(weight_name, param_name)

1636 + param_name = weight_name.removeprefix(f"{self.layer_name}.")

1637 + param = getattr(self, param_name)

1638 + success = self.weight_loader(

1639 + param=param,

1640 + loaded_weight=loaded_weight,

1641 + weight_name=weight_name,

1642 + shard_id=shard_id,

1643 + expert_id=expert_id,

1644 + return_success=True,

1645 + )

1646 + if success:

1647 + logger.debug("Loaded %s for expert %d into %s", param_name,

1648 + expert_id, self.layer_name)

1649 + yield param_name

1650 +

1651 def get_expert_weights(self) -> Iterable[torch.Tensor]:

1652 weights = list(self.named_parameters())

1653 assert all(weight.is_contiguous() for _, weight in weights)

vllm/model_executor/models/registry.py

+6 -2

@@ -307,10 +307,14 @@ _TRANSFORMERS_SUPPORTED_MODELS = {

307 }

308

309 _TRANSFORMERS_BACKEND_MODELS = {

310 - "TransformersEmbeddingModel": ("transformers_pooling", "TransformersEmbeddingModel"), # noqa: E501

311 - "TransformersForSequenceClassification": ("transformers_pooling", "TransformersForSequenceClassification"), # noqa: E501

310 "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),

311 "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501

312 + "TransformersMoEForCausalLM": ("transformers_moe", "TransformersMoEForCausalLM"), # noqa: E501

313 + "TransformersMoEForMultimodalLM": ("transformers_moe", "TransformersMoEForMultimodalLM"), # noqa: E501

314 + "TransformersEmbeddingModel": ("transformers_pooling", "TransformersEmbeddingModel"), # noqa: E501

315 + "TransformersForSequenceClassification": ("transformers_pooling", "TransformersForSequenceClassification"), # noqa: E501

316 + "TransformersMoEForSequenceClassification": ("transformers_pooling", "TransformersMoEForSequenceClassification"), # noqa: E501

317 + "TransformersMoEEmbeddingModel": ("transformers_pooling", "TransformersMoEEmbeddingModel"), # noqa: E501

318 }

319 # yapf: enable

320

vllm/model_executor/models/transformers.py

+107 -54

@@ -22,6 +22,8 @@ from typing import Literal, Optional, Union

23 import regex as re

24 import torch

25 +import transformers

26 +from packaging.version import Version

27 from torch import nn

28 from transformers import (AutoModel, BatchFeature, PretrainedConfig,

29 PreTrainedModel)

@@ -35,6 +37,7 @@ from vllm.config.utils import getattr_iter

37 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

38 from vllm.distributed.utils import get_pp_indices

39 from vllm.logger import init_logger

40 +from vllm.model_executor.layers.layernorm import RMSNorm

41 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

42 ReplicatedLinear,

43 RowParallelLinear)

@@ -121,10 +124,14 @@ def can_enable_torch_compile(vllm_config: VllmConfig) -> bool:

124 return enable

125

126

127 +Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep",

128 + "replicate"]

129 +

130 +

131 def replace_linear_class(

132 linear: nn.Linear,

126 - style: Literal["colwise", "rowwise"],

127 - quant_config: QuantizationConfig,

133 + style: Style = "replicate",

134 + quant_config: Optional[QuantizationConfig] = None,

135 *,

136 prefix: str = "",

137 ) -> Union[ColumnParallelLinear, RowParallelLinear, ReplicatedLinear]:

@@ -132,11 +139,11 @@ def replace_linear_class(

139 Replace nn.Linear with one of vLLM's tensor parallel linear classes.

140

141 Args:

135 - linear (nn.Linear): `nn.Linear` to be replaced.

136 - style (str): Tensor parallel style of the new linear, e.g. "colwise".

137 - quant_config (QuantConfig): Quantization config for the new linear.

142 + linear: `nn.Linear` to be replaced.

143 + style: Tensor parallel style of the new linear, e.g. "colwise".

144 + quant_config: Quantization config for the new linear.

145 Returns:

139 - Union[ColumnParallelLinear, RowParallelLinear]: The new linear.

146 + The new linear.

147 """

148

149 if not isinstance(style, str):

@@ -166,6 +173,31 @@ def replace_linear_class(

173 )

174

175

176 +def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:

177 + """Replace a Transformers RMSNorm with vLLM's RMSNorm.

178 +

179 + This method assumes:

180 + - Weight is stored as `weight`.

181 + - Epsilon is stored as `eps` or `variance_epsilon`.

182 + - `with_scale` indicates whether the layer has a weight (Gemma3n only).

183 + - `var_hidden_size` is only ever used for Intern vision encoder in vLLM

184 + and Transformers doesn't appear to have the same concept.

185 + """

186 + kwargs = {

187 + "hidden_size": hidden_size,

188 + "eps": getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6),

189 + "has_weight": getattr(rms_norm, "with_scale", True)

190 + }

191 + if (weight := getattr(rms_norm, "weight", None)) is not None:

192 + # If weight is a Parameter, get its data tensor

193 + weight = getattr(weight, "data", weight)

194 + kwargs["dtype"] = weight.dtype

195 + else:

196 + # No weight, fall back to weightless RMSNorm

197 + kwargs["has_weight"] = False

198 + return RMSNorm(**kwargs)

199 +

200 +

201 # Copied from `accelerate`

202 @contextmanager

203 def init_on_device_without_buffers(device: torch.device):

@@ -463,9 +495,15 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

495 self.ignore_unexpected_suffixes: list[str] = []

496 """Ignore unexpected weights whose qualname ends with these suffixes."""

497

466 - # Skip loading extra bias for GPTQ models.

467 - if self.quant_config and "gptq" in self.quant_config.get_name():

468 - self.ignore_unexpected_suffixes.append(".bias")

498 + if self.quant_config:

499 + quant_method_name = self.quant_config.get_name()

500 + # Check for unsupported quantization methods.

501 + if quant_method_name == "mxfp4":

502 + raise NotImplementedError("Transformers backend does not "

503 + "support MXFP4 quantization yet.")

504 + # Skip loading extra bias for GPTQ models.

505 + if "gptq" in quant_method_name:

506 + self.ignore_unexpected_suffixes.append(".bias")

507

508 # Set correct attn and init on "meta" to delay allocating GPU tensors

509 # TODO: @raushan, use the public `model.set_attn_implementation()`

@@ -478,8 +516,12 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

516 trust_remote_code=self.model_config.trust_remote_code,

517 )

518

519 + # Remove layers not on this pipeline parallel rank

520 self.pipeline_parallel()

482 - self.tensor_parallel()

521 + # Substitute remaining layers with vLLM's layers as needed

522 + self.recursive_replace()

523 + # Create attention instances for KV cache allocation

524 + self.attention_instances = self.create_attention_instances()

525

526 # Input embeddings

527 if not isinstance(self.model.get_input_embeddings(), PPMissingLayer):

@@ -494,12 +536,10 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

536 quant_config=self.quant_config,

537 ))

538

497 - # Attention layers

498 - self.attention_instances = self.create_attention_instances()

499 -

539 # Initialize any parameters that have not had their modules replaced

540 self.init_parameters(self.model)

541

542 + # Pipeline parallel intermediate tensors

543 self.make_empty_intermediate_tensors = (

544 make_empty_intermediate_tensors_factory(

545 ["hidden_states"], self.text_config.hidden_size))

@@ -558,56 +598,53 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

598 if not self.pp_group.is_last_rank:

599 setattr(self.model, name, PPMissingLayer())

600

561 - def tensor_parallel(self):

562 - """

563 - Apply the model's tensor parallelization plan.

564 - Currently only supports linear layers.

601 + def recursive_replace(self):

602 + """Recursively replace modules in the model as needed.

603 +

604 + Currently, this replaces:

605 +

606 + - `nn.Linear` with vLLM's tensor parallel linear classes

607 + - `*RMSNorm` with vLLM's `RMSNorm`

608 """

566 - # Look for tp plans in all of the PreTrainedModels found in self.model

567 - is_pretrained_model = lambda m: isinstance(m, PreTrainedModel)

568 - supports_tp_plan = lambda m: m.config.base_model_tp_plan is not None

569 - pretrained_models = filter(is_pretrained_model, self.model.modules())

570 - models_with_tp_plan = filter(supports_tp_plan, pretrained_models)

609 + tp_plan = self.model.tp_plan

610

572 - if not any(models_with_tp_plan) and self.tp_size > 1:

611 + if not tp_plan and self.tp_size > 1:

612 tip = get_feature_request_tip(self.model_config.model,

613 self.model_config.trust_remote_code)

614 raise ValueError(

615 f"{type(self.model)} does not support tensor parallel. {tip}")

616

578 - def _tensor_parallel(module: nn.Module, prefix: str, tp_plan=None):

579 - tp_plan = tp_plan or {}

580 -

581 - # If the current module is a PreTrainedModel, set the tp_plan for

582 - # all of its children

583 - if isinstance(module, PreTrainedModel):

584 - tp_plan = module.config.base_model_tp_plan or {}

585 - tp_plan = {

586 - maybe_prefix(prefix, k): v

587 - for k, v in tp_plan.items()

588 - }

589 -

590 - # Some weight loaders expect linear layers to inherit from vLLM's

591 - # LinearBase class, so we set a default style which causes any

592 - # unspecified linear layers to be replaced with ReplicatedLinear

617 + # Prefix the patterns because we always start from `self.model`

618 + tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}

619 +

620 + def _recursive_replace(module: nn.Module, prefix: str):

621 for child_name, child_module in module.named_children():

622 + new_module = child_module

623 qual_name = maybe_prefix(prefix, child_name)

624 if isinstance(child_module, nn.Linear):

625 generator = (p for p in tp_plan if re.match(p, qual_name))

626 pattern = next(generator, None)

627 + # Some weight loaders expect all linear layers to inherit

628 + # LinearBase, so we set a default style which causes any

629 + # unspecified layers to be replaced with ReplicatedLinear

630 style = tp_plan.get(pattern, "replicate")

631 new_module = replace_linear_class(child_module,

632 style,

633 self.quant_config,

634 prefix=qual_name)

635 + # TODO(hmellor): Enable RMSNorm replacement once we have a way

636 + # to choose RMSNorm vs GemmaRMSNorm

637 + # elif child_module.__class__.__name__.endswith("RMSNorm"):

638 + # new_module = replace_rms_norm_class(

639 + # child_module, self.config.hidden_size)

640 + else:

641 + _recursive_replace(child_module, prefix=qual_name)

642 +

643 + if new_module is not child_module:

644 setattr(module, child_name, new_module)

645 log_replacement(qual_name, child_module, new_module)

605 - else:

606 - _tensor_parallel(child_module,

607 - prefix=qual_name,

608 - tp_plan=tp_plan)

646

610 - _tensor_parallel(self.model, prefix="model")

647 + _recursive_replace(self.model, prefix="model")

648

649 def create_attention_instances(

650 self,

@@ -657,15 +694,21 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

694 self.model: PreTrainedModel = AutoModel.from_config(...)

695 ```

696 """

660 - for name, param in module.named_parameters(recurse=False):

661 - if param.device == torch.device("meta"):

662 - new_param = nn.Parameter(

663 - torch.empty_like(param.data,

664 - dtype=dtype or self.model_config.dtype,

665 - device=self.device_config.device))

666 - setattr(module, name, new_param)

667 - for child in module.children():

668 - self.init_parameters(child, dtype)

697 +

698 + def _init_parameters(module: nn.Module, dtype: Optional[torch.dtype]):

699 + for name, param in module.named_parameters(recurse=False):

700 + if param.device == torch.device("meta"):

701 + new_param = nn.Parameter(

702 + torch.empty_like(

703 + param.data,

704 + dtype=dtype or self.model_config.dtype,

705 + device=self.device_config.device,

706 + ))

707 + setattr(module, name, new_param)

708 + for child in module.children():

709 + _init_parameters(child, dtype)

710 +

711 + _init_parameters(module, dtype)

712

713 def forward(

714 self,

@@ -702,8 +745,10 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

745

746 return hidden_states

747

705 - def load_weights(self, weights: Iterable[tuple[str,

706 - torch.Tensor]]) -> set[str]:

748 + def load_weights(

749 + self,

750 + weights: Iterable[tuple[str, torch.Tensor]],

751 + ) -> set[str]:

752 loader = AutoWeightsLoader(

753 self,

754 skip_prefixes=self.skip_prefixes,

@@ -713,6 +758,14 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

758 )

759 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

760

761 + def check_version(self, min_version: str, feature: str):

762 + installed = Version(transformers.__version__)

763 + required = Version(min_version)

764 + if installed < required:

765 + raise ImportError(

766 + f"Transformers backend requires transformers>={required} "

767 + f"for {feature}, but got {installed}")

768 +

769

770 @support_torch_compile(enable_if=can_enable_torch_compile)

771 class TransformersForCausalLM(TransformersBase):

vllm/model_executor/models/transformers_moe.py new

+268

@@ -0,0 +1,268 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

5 +#

6 +# Licensed under the Apache License, Version 2.0 (the "License");

7 +# you may not use this file except in compliance with the License.

8 +# You may obtain a copy of the License at

9 +#

10 +# http://www.apache.org/licenses/LICENSE-2.0

11 +#

12 +# Unless required by applicable law or agreed to in writing, software

13 +# distributed under the License is distributed on an "AS IS" BASIS,

14 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

15 +# See the License for the specific language governing permissions and

16 +# limitations under the License.

17 +"""Wrapper around `transformers` MoE models."""

18 +from typing import Any

19 +

20 +import torch

21 +import torch.nn as nn

22 +

23 +from vllm.compilation.decorators import support_torch_compile

24 +from vllm.config.utils import getattr_iter

25 +from vllm.forward_context import ForwardContext, get_forward_context

26 +from vllm.model_executor.custom_op import CustomOp

27 +from vllm.model_executor.layers.fused_moe import FusedMoE

28 +from vllm.platforms import current_platform

29 +from vllm.utils import direct_register_custom_op

30 +

31 +from .transformers import (TransformersBase, TransformersForCausalLM,

32 + TransformersForMultimodalLM,

33 + can_enable_torch_compile, log_replacement)

34 +from .utils import maybe_prefix

35 +

36 +

37 +@CustomOp.register("transformers_fused_moe")

38 +class TransformersFusedMoE(FusedMoE):

39 + """Custom FusedMoE for the Transformers backend."""

40 +

41 + def __init__(self, *args, **kwargs):

42 + super().__init__(*args, **kwargs)

43 + self._top_k_index: torch.Tensor = None

44 +

45 + def custom_routing_function(hidden_states, gating_output, topk,

46 + renormalize):

47 + """Return `top_k_weights` from `gating_output` and the

48 + `top_k_index` we stored in the layer earlier."""

49 + return gating_output, self._top_k_index

50 +

51 + self.custom_routing_function = custom_routing_function

52 +

53 + def forward(self, hidden_states: torch.Tensor, top_k_index: torch.Tensor,

54 + top_k_weights: torch.Tensor, **kwargs: Any) -> torch.Tensor:

55 + """In Transformers `experts.forward` will have this signature.

56 +

57 + We discard any extra kwargs because we cannot use them here."""

58 + return torch.ops.vllm.transformers_moe_forward(hidden_states,

59 + top_k_index,

60 + top_k_weights,

61 + self.layer_name)

62 +

63 +

64 +def transformers_moe_forward(hidden_states: torch.Tensor,

65 + top_k_index: torch.Tensor,

66 + top_k_weights: torch.Tensor,

67 + layer_name: str) -> torch.Tensor:

68 + """Store the `top_k_index` in the layer and call the actual forward."""

69 + forward_context: ForwardContext = get_forward_context()

70 + self = forward_context.no_compile_layers[layer_name]

71 + self._top_k_index = top_k_index

72 + # Clone hidden_states because it will be mutated in-place in FusedMoE

73 + return self.forward_impl(hidden_states.clone(), top_k_weights)

74 +

75 +

76 +def transformers_moe_forward_fake(hidden_states: torch.Tensor,

77 + top_k_index: torch.Tensor,

78 + top_k_weights: torch.Tensor,

79 + layer_name: str) -> torch.Tensor:

80 + return torch.empty_like(hidden_states)

81 +

82 +

83 +direct_register_custom_op(

84 + op_name="transformers_moe_forward",

85 + op_func=transformers_moe_forward,

86 + mutates_args=["hidden_states"],

87 + fake_impl=transformers_moe_forward_fake,

88 + dispatch_key=current_platform.dispatch_key,

89 + tags=(torch.Tag.needs_fixed_stride_order, ),

90 +)

91 +

92 +

93 +class TransformersMoEBase(TransformersBase):

94 +

95 + def __init__(self, *, vllm_config, prefix=""):

96 + self.check_version("4.57.0.dev0", "MoE models support")

97 + super().__init__(vllm_config=vllm_config, prefix=prefix)

98 +

99 + if self.parallel_config.enable_expert_parallel:

100 + raise NotImplementedError(

101 + "Transformers backend does not support expert parallel yet.")

102 + if self.parallel_config.enable_eplb:

103 + raise NotImplementedError(

104 + "Transformers backend does not support expert parallel load "

105 + "balancing yet.")

106 +

107 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

108 + """

109 + Params for weights, fp8 weight scales, fp8 activation scales

110 + (param_name, weight_name, expert_id, shard_id)

111 + """

112 + ckpt_names = [

113 + # (ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name)

114 + ("gate_proj", "down_proj", "up_proj"), # Most common MoE style

115 + ("w1", "w2", "w3"), # Granite, Mixtral, Phi MoE style

116 + ("linear", "linear_1", "linear_v"), # Grok1 style

117 + ]

118 + expert_mapping = []

119 + for gate_proj, down_proj, up_proj in ckpt_names:

120 + expert_mapping.extend(

121 + FusedMoE.make_expert_params_mapping(

122 + ckpt_gate_proj_name=gate_proj,

123 + ckpt_down_proj_name=down_proj,

124 + ckpt_up_proj_name=up_proj,

125 + num_experts=self.model_config.get_num_experts(),

126 + num_redundant_experts=0, # TODO: enable EPLB

127 + ))

128 + return expert_mapping

129 +

130 + def recursive_replace(self):

131 + """Initialize the MoE layers."""

132 + text_config = self.text_config

133 +

134 + # Positional arguments

135 + num_experts = self.model_config.get_num_experts()

136 + top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"],

137 + None)

138 + assert top_k is not None

139 + hidden_size = text_config.hidden_size

140 + intermediate_size = getattr_iter(

141 + text_config, ["moe_intermediate_size", "intermediate_size"], None)

142 + assert intermediate_size is not None

143 +

144 + # If there are shared experts, the results are

145 + # reduced after mlp.forward() not inside FusedMoE

146 + num_experts_shared = getattr_iter(text_config, [

147 + "num_experts_shared", "n_shared_experts", "moe_num_shared_experts"

148 + ], 0)

149 + reduce_results = num_experts_shared == 0

150 +

151 + def add_all_reduce(mlp: nn.Module):

152 + """Adds an all-reduce to the output of `mlp.forward()`."""

153 +

154 + class MLPWithAllReduce(mlp.__class__):

155 +

156 + def forward(self, *args, **kwargs):

157 + output = super().forward(*args, **kwargs)

158 + return self.experts.maybe_all_reduce_tensor_model_parallel(

159 + output)

160 +

161 + mlp.__class__ = MLPWithAllReduce

162 +

163 + # Unused kwargs since we use custom_routing_function:

164 + # - `scoring_func` and `e_score_correction_bias` only used for grouped

165 + # topk routing inside vLLM and are non-trivial to infer

166 + # and hard code `use_grouped_topk=False`

167 + # - `renormalize` passed anyway because it's easy to infer

168 + # - `num_expert_group` and `topk_group` used for inferring expert

169 + # placement strategy in FusedMoE

170 + # - `apply_router_weight_on_input` is already applied in Transformers

171 + renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)

172 + num_expert_group = getattr(text_config, "n_group", None)

173 + topk_group = getattr(text_config, "topk_group", None)

174 +

175 + # MoE activation function

176 + activation = "silu"

177 + wrapped_arch = self.config.architectures[0].lower()

178 + if "gptoss" in wrapped_arch:

179 + activation = "swigluoai"

180 + elif "grok1" in wrapped_arch:

181 + activation = "gelu"

182 +

183 + # Expert mapping for `AutoWeightsLoader`

184 + expert_mapping = self.get_expert_mapping()

185 +

186 + # Configs

187 + parallel_config = self.parallel_config

188 + eplb_config = parallel_config.eplb_config

189 +

190 + # Expert parallel load balancing kwargs

191 + enable_eplb = parallel_config.enable_eplb

192 + num_redundant_experts = eplb_config.num_redundant_experts

193 +

194 + # Recursively fuse MoE layers

195 + def _recursive_replace(module: nn.Module, prefix: str):

196 + for child_name, child_module in module.named_children():

197 + qual_name = maybe_prefix(prefix, child_name)

198 + if (child_name == "experts"

199 + and isinstance(child_module, nn.ModuleList)):

200 + # Alias for readability

201 + mlp = module

202 + experts = child_module

203 + # Do the experts have biases

204 + has_bias = False

205 + for experts_param_name, _ in experts.named_parameters():

206 + if "bias" in experts_param_name:

207 + has_bias = True

208 + break

209 + # Double check there are no shared experts

210 + nonlocal reduce_results

211 + if reduce_results:

212 + for mlp_param_name, _ in mlp.named_parameters():

213 + if "shared_expert" in mlp_param_name:

214 + reduce_results = False

215 + break

216 + # Replace experts module with FusedMoE

217 + fused_experts = TransformersFusedMoE(

218 + num_experts=num_experts,

219 + top_k=top_k,

220 + hidden_size=hidden_size,

221 + intermediate_size=intermediate_size,

222 + reduce_results=reduce_results,

223 + renormalize=renormalize,

224 + # Hard coded because topk happens in Transformers

225 + use_grouped_topk=False,

226 + num_expert_group=num_expert_group,

227 + topk_group=topk_group,

228 + quant_config=self.quant_config,

229 + prefix=qual_name,

230 + activation=activation,

231 + enable_eplb=enable_eplb,

232 + num_redundant_experts=num_redundant_experts,

233 + has_bias=has_bias,

234 + expert_mapping=expert_mapping,

235 + )

236 + mlp.experts = fused_experts

237 + log_replacement(qual_name, experts, fused_experts)

238 + # If results are not all-reduced in FusedMoE, ensure they

239 + # are all-reduced at the end of mlp.forward() if tensor

240 + # parallel or expert parallel is enabled

241 + if not reduce_results and (fused_experts.tp_size > 1

242 + or fused_experts.ep_size > 1):

243 + add_all_reduce(mlp)

244 + else:

245 + _recursive_replace(child_module, prefix=qual_name)

246 +

247 + _recursive_replace(self.model, prefix="model")

248 + # Continue with the replacement of layers in TransformersBase

249 + super().recursive_replace()

250 +

251 +

252 +@support_torch_compile(enable_if=can_enable_torch_compile)

253 +class TransformersMoEForCausalLM(TransformersMoEBase, TransformersForCausalLM):

254 + pass

255 +

256 +

257 +@support_torch_compile(

258 + # set `positions` to last dim to support Qwen-mrope

259 + dynamic_arg_dims={

260 + "input_ids": 0,

261 + "positions": -1,

262 + "intermediate_tensors": 0,

263 + "inputs_embeds": 0,

264 + },

265 + enable_if=can_enable_torch_compile)

266 +class TransformersMoEForMultimodalLM(TransformersMoEForCausalLM,

267 + TransformersForMultimodalLM):

268 + pass

vllm/model_executor/models/transformers_pooling.py

+18 -10

@@ -20,7 +20,7 @@ from typing import Optional, Union

20 import torch

21 from transformers import AutoModelForSequenceClassification

23 -from vllm.attention import AttentionType

23 +from vllm.attention import Attention, AttentionType

24 from vllm.compilation.decorators import support_torch_compile

25 from vllm.config import VllmConfig

26 from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,

@@ -29,6 +29,7 @@ from vllm.sequence import IntermediateTensors

30 from .interfaces_base import VllmModelForPooling

31 from .transformers import TransformersBase, can_enable_torch_compile

32 +from .transformers_moe import TransformersMoEBase

33 from .utils import WeightsMapper

@@ -79,7 +80,9 @@ class TransformersPoolingBase(TransformersBase, VllmModelForPooling):

80 self.padding_idx = self.text_config.pad_token_id

82 def create_attention_instances(

82 - self, attn_type: AttentionType = AttentionType.DECODER):

83 + self,

84 + attn_type: AttentionType = AttentionType.DECODER

85 + ) -> dict[int, Attention]:

86 # TODO(hmellor): Better way to detect encoder models

87 # In encoder models, the attention layers will have `is_causal=False`

88 is_encoder = lambda m: not getattr(m, "is_causal", True)

@@ -90,14 +93,7 @@ class TransformersPoolingBase(TransformersBase, VllmModelForPooling):

94 # Check minimum transformers version for encoder models support

95 if attn_type == AttentionType.ENCODER_ONLY:

93 - import transformers

94 - from packaging.version import Version

95 - installed = Version(transformers.__version__)

96 - required = Version("4.57.0.dev0")

97 - if installed < required:

98 - raise ValueError(

99 - "Encoder models with the Transformers backend require "

100 - f"transformers>={required}, but got {installed}")

96 + self.check_version("4.57.0.dev0", "encoder models support")

98 return super().create_attention_instances(attn_type)

@@ -198,3 +194,15 @@ class TransformersForSequenceClassification(TransformersPoolingBase):

194 vllm_config.model_config),

195 ),

196 })

197 +

198 +

199 +@support_torch_compile(enable_if=can_enable_torch_compile)

200 +class TransformersMoEEmbeddingModel(TransformersMoEBase,

201 + TransformersEmbeddingModel):

202 + pass

203 +

204 +

205 +@support_torch_compile(enable_if=can_enable_torch_compile)

206 +class TransformersMoEForSequenceClassification(

207 + TransformersMoEBase, TransformersForSequenceClassification):

208 + pass