@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 150 lines 11.1 KB
Rendered Raw
1 # Basic Model
2
3 This guide walks you through the steps to implement a basic vLLM model.
4
5 ## 1. Bring your model code
6
7 First, clone the PyTorch model code from the source repository.
8 For instance, vLLM's [OPT model](../../../vllm/model_executor/models/opt.py) was adapted from
9 HuggingFace's [modeling_opt.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py) file.
10
11 !!! warning
12 Make sure to review and adhere to the original code's copyright and licensing terms!
13
14 ## 2. Make your code compatible with vLLM
15
16 To ensure compatibility with vLLM, your model must meet the following requirements:
17
18 ### Initialization Code
19
20 All vLLM modules within the model must include a `prefix` argument in their constructor. This `prefix` is typically the full name of the module in the model's state dictionary and is crucial for:
21
22 - Runtime support: vLLM's attention operators are registered in a model's state by their full names. Each attention operator must have a unique prefix as its layer name to avoid conflicts.
23 - Non-uniform quantization support: A quantized checkpoint can selectively quantize certain layers while keeping others in full precision. By providing the `prefix` during initialization, vLLM can match the current layer's `prefix` with the quantization configuration to determine if the layer should be initialized in quantized mode.
24
25 The initialization code should look like this:
26
27 ??? code
28
29 ```python
30 from torch import nn
31 from vllm.config import VllmConfig
32 from vllm.model_executor.layers.attention import Attention
33
34 class MyAttention(nn.Module):
35 def __init__(self, vllm_config: VllmConfig, prefix: str):
36 super().__init__()
37 self.attn = Attention(prefix=f"{prefix}.attn")
38
39 class MyDecoderLayer(nn.Module):
40 def __init__(self, vllm_config: VllmConfig, prefix: str):
41 super().__init__()
42 self.self_attn = MyAttention(prefix=f"{prefix}.self_attn")
43
44 class MyModel(nn.Module):
45 def __init__(self, vllm_config: VllmConfig, prefix: str):
46 super().__init__()
47 self.layers = nn.ModuleList(
48 [MyDecoderLayer(vllm_config, prefix=f"{prefix}.layers.{i}") for i in range(vllm_config.model_config.hf_config.num_hidden_layers)]
49 )
50
51 class MyModelForCausalLM(nn.Module):
52 def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
53 super().__init__()
54 self.model = MyModel(vllm_config, prefix=f"{prefix}.model")
55 ```
56
57 ### Computation Code
58
59 - Add a `embed_input_ids` method inside `MyModel` module that returns the text embeddings given `input_ids`. This is equivalent to directly calling the text embedding layer, but provides a unified interface in case `MyModel` is used within a composite multimodal model.
60
61 ```python
62 class MyModel(nn.Module):
63 ...
64
65 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
66 ...
67 ```
68
69 - Rewrite the [forward][torch.nn.Module.forward] method of your model to remove any unnecessary code, such as training-specific code. Modify the input parameters to treat `input_ids` and `positions` as flattened tensors with a single batch size dimension, without a max-sequence length dimension.
70
71 ```python
72 def forward(
73 self,
74 input_ids: torch.Tensor | None,
75 positions: torch.Tensor,
76 intermediate_tensors: IntermediateTensors | None = None,
77 inputs_embeds: torch.Tensor | None = None,
78 ) -> torch.Tensor:
79 ...
80 ```
81
82 !!! note
83 Currently, vLLM supports the basic multi-head attention mechanism and its variant with rotary positional embeddings.
84 If your model employs a different attention mechanism, you will need to implement a new attention layer in vLLM.
85
86 For reference, check out our [Llama implementation](../../../vllm/model_executor/models/llama.py). vLLM already supports a large number of models. It is recommended to find a model similar to yours and adapt it to your model's architecture. Check out [vllm/model_executor/models](../../../vllm/model_executor/models) for more examples.
87
88 Adapting model-specific structure or boilerplate can keep architectures independent. Reuse shared layer implementations, such as Attention and MoE, where their contracts fit the model.
89
90 ## 3. (Optional) Implement tensor parallelism and quantization support
91
92 If your model is too large to fit into a single GPU, you can use tensor parallelism to manage it.
93 To do this, substitute your model's linear and embedding layers with their tensor-parallel versions.
94 For the embedding layer, you can simply replace [torch.nn.Embedding][] with `VocabParallelEmbedding`. For the output LM head, you can use `ParallelLMHead`.
95 When it comes to the linear layers, we provide the following options to parallelize them:
96
97 - `ReplicatedLinear`: Replicates the inputs and weights across multiple GPUs. No memory saving.
98 - `RowParallelLinear`: The input tensor is partitioned along the hidden dimension. The weight matrix is partitioned along the rows (input dimension). An *all-reduce* operation is performed after the matrix multiplication to reduce the results. Typically used for the second FFN layer and the output linear transformation of the attention layer.
99 - `ColumnParallelLinear`: The input tensor is replicated. The weight matrix is partitioned along the columns (output dimension). The result is partitioned along the column dimension. Typically used for the first FFN layer and the separated QKV transformation of the attention layer in the original Transformer.
100 - `MergedColumnParallelLinear`: Column-parallel linear that merges multiple `ColumnParallelLinear` operators. Typically used for the first FFN layer with weighted activation functions (e.g., SiLU). This class handles the sharded weight loading logic of multiple weight matrices.
101 - `QKVParallelLinear`: Parallel linear layer for the query, key, and value projections of the multi-head and grouped-query attention mechanisms. When number of key/value heads are less than the world size, this class replicates the key/value heads properly. This class handles the weight loading and replication of the weight matrices.
102
103 Note that all the linear layers above take `linear_method` as an input. vLLM will set this parameter according to different quantization schemes to support weight quantization.
104
105 ## 4. Implement the weight loading logic
106
107 You now need to implement the `load_weights` method in your `*ForCausalLM` class.
108 This method should load the weights from the HuggingFace's checkpoint file and assign them to the corresponding layers in your model. Specifically, for `MergedColumnParallelLinear` and `QKVParallelLinear` layers, if the original model has separated weight matrices, you need to load the different parts separately.
109
110 ## 5. Register your model
111
112 See [this page](registration.md) for instructions on how to register your new model to be used by vLLM.
113
114 ## Frequently Asked Questions
115
116 ### How to support models with interleaving sliding windows?
117
118 To support a model with interleaving sliding windows, we need to take care of the following details:
119
120 - Make sure the model's `config.json` contains `layer_types`.
121 - In the modeling code, parse the correct sliding window value for every layer, and pass it to the attention layer's `per_layer_sliding_window` argument. For reference, check [this line](https://github.com/vllm-project/vllm/blob/996357e4808ca5eab97d4c97c7d25b3073f46aab/vllm/model_executor/models/llama.py#L171).
122
123 With these two steps, interleaved sliding windows should work with the model.
124
125 ### How to support models that use Mamba?
126
127 We consider 3 different scenarios:
128
129 1. Models that use Mamba layers (either Mamba-1 or Mamba-2) but do not use attention layers.
130 2. Models that combine Mamba layers (either Mamba-1 or Mamba-2) together with attention layers.
131 3. Models that combine Mamba-like mechanisms (e.g., Linear Attention, ShortConv) together with attention layers.
132
133 For case (1), we recommend looking at the implementation of [`MambaForCausalLM`](../../../vllm/model_executor/models/mamba.py) (for Mamba-1) or [`Mamba2ForCausalLM`](../../../vllm/model_executor/models/mamba2.py) (for Mamba-2) as a reference.
134 The model should inherit protocol `IsAttentionFree` and also implement class methods `get_mamba_state_dtype_from_config` and `get_mamba_state_shape_from_config` to calculate the state shapes and data types from the config.
135 For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
136 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
137
138 For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
139 These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
140
141 For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
142 Please follow the same guidelines as case (2) for implementing these models.
143 We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
144 For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
145 It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
146 Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
147 It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
148 Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
149 Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
150 The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.