[Core][Model][Frontend] Model architecture plugins

[NadavShmayo]

Solves #7124 (also relevant for #7131).

This pull request implements support for external model architecture plugins.
As discussed in #7131, vLLM should support both "general purpose" plugins which allow arbitrary modifications, but in addition it should allow simple plugin interfaces for parts in the vLLM architecture which might change a lot.

Using this pull request we can simply create a separate package implementing a model architecture, and install it in the same environment as vLLM.
If the model architecture plugin overrides an existing architecture you can use the --model-architecture-override option to specify the name of the new architecture to use.

Code for a mock OPT architecture implementation

from vllm import ModelRegistry
from vllm.plugins.model_plugin import ModelArchitectureBase, ModelPlugin
from vllm.model_executor.models.opt import OPTForCausalLM


class MyOpt(OPTForCausalLM):
    def __init__(self, config, cache_config, quant_config):
        super().__init__(config, cache_config, quant_config)
    def forward(self, input_ids, positions, kv_caches, attn_metadata, intermediate_tensors=None):
        res = OPTForCausalLM.forward(self, input_ids, positions, kv_caches, attn_metadata, intermediate_tensors)
        res[0] = 1000
        return res


MODEL_PLUGIN = ModelPlugin(architecture_name='MyOpt', implementation_cls=MyOpt)

Setup.py file

from setuptools import setup

setup(name='test_plugin',
      version='0.1',
      description='Example model architecture plugin for vLLM',
      install_requires=[
      ],
      entry_points={
            'vllm.model_architectures': ['test_plugin=test_plugin.main:MODEL_PLUGIN'],
      }
      )

Todo list

• Implement basic logic
• Add test
• Add documentation

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

Will take a look after #7426

[DarkLight1337]

I think the model class also needs to implement compute_logits and sample in order to be used in the model runner. So the full implementation would be like

from typing import Protocol, runtime_checkable
from typing_extensions import TypeVar

# Currently, T = torch.Tensor for all models except Medusa, which has T = List[torch.Tensor]
T = TypeVar("T", default=torch.Tensor)

@runtime_checkable
class SupportsVllm(Protocol[T]):
    def __init__(
        self,
        config: PreTrainedConfig,
        *,
        # In vllm.model_executor.model_loader, the model is always constructed by passing these via kwargs
        cache_config: Optional[CacheConfig],
        quant_config: Optional[QuantizationConfig],
    ) -> None:
        ...

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> T:
        ...

    def compute_logits(
        self,
        hidden_states: T,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[T]:
        """Return `None` if rank > 0."""
        ...

    def sample(
        self,
        logits: T,
        sampling_metadata: SamplingMetadata,
    ) -> SamplerOutput:
        """Only called on rank 0."""
        ...

[youkaichao]

I feel this is quite complicated, but actually we just need one line ModelRegistry.register_model .

I would prefer it remains a general plugin.

I think we need to come up with a new plugin type, only when it is very difficult to achieve through existing one.

[NadavShmayo]

I feel this is quite complicated, but actually we just need one line ModelRegistry.register_model .

I would prefer it remains a general plugin.

I think we need to come up with a new plugin type, only when it is very difficult to achieve through existing one.

I can see what you are saying, I think it would be acceptable to enable this functionality with general plugins, but I don't see much downside to supporting this specific type of plugin, as it should be used quite frequently.

Regardless of how we support this, I believe we should at least add the --model-architecture-override argument, so you don't have to actually override the config.json file to use a plugin.

[youkaichao]

but I don't see much downside to supporting this specific type of plugin

I will only accept PRs that bring evident benefit. No downside is not a good pitch. I try to avoid the codebase being too swollen.

I believe we should at least add the --model-architecture-override argument, so you don't have to actually override the config.json file to use a plugin.

I don't understand this. With ModelRegistry.register_model, you can already override existing architectures.

[NadavShmayo]

but I don't see much downside to supporting this specific type of plugin

I will only accept PRs that bring evident benefit. No downside is not a good pitch. I try to avoid the codebase being too swollen.

I believe we should at least add the --model-architecture-override argument, so you don't have to actually override the config.json file to use a plugin.

I don't understand this. With ModelRegistry.register_model, you can already override existing architectures.

I didn't realize register_model could override existing implementations.
In that case, this PR really isn't required.

In this case, it's really simple to create a general purpose plugin to achieve the goal, so it's not necessary to create a specific plugin type. But I do believe in most cases, unless very simple to implement, we should lean towards specific plugin types to make the plugin system more robust.

Do you agree? Furthermore, can you think of a plugin type that could already be implemented in the existing state of vLLM?