deepspeedai--deepspeed/deepspeed/inference/v2/model_implementations/AddingAModel.md

Adding Support for a New Model in DeepSpeed Inference V2

Adding supoprt for a new model in DeepSpeed Inference requires developing three related components:

• Containers: These describe the parameters contained in the model
• Model implementation: How should the model be computed.
• Policy: The map for adding parameters to your containers and creating the model implementation.

In this tutorial, we will assume that you'd like to use a relatively traditionally styled Transformer model and will be able to inherit from DSTransformerModelBase and can take advantage of the utilities that provides.

Defining Your Containers

A container is the bridge between the original model's parameters and how to transform them to serve them for inference. For a model implementation, there are two primary kinds of containers: transformer containers and non-transformer containers. A transformer container consists of the parameters for a single Transformer layer in the model. So this includes your traditional parameters like the projections for the fully connected network, or query-key-value projections. The non-transformer container will contain basically everything else! However, before defining these containers, we need to understand how to define an individual parameter.

In DeepSpeed inference, the original model parameters are populated into the model and mapped as dependencies to a parameter. A Parameter has two primary components: its dependencies and its finalize method. Let's do an example. In Llama models, the native format is for the query, key, and value projections to be performed independently. However, we can achieve higher throughput by fusing them into a single larger projection. We can define this fusion with a parameter:

from deepspeed.inference.module_implementations.parameter_base import ParameterBase

class UnfusedQKVParameter(ParameterBase):
    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor

    def finalize(self) -> torch.Tensor:
        fused_param = torch.cat([self.query, self.key, self.value], dim=0)
        return self.inference_model.transform_qkv_param(fused_param)

Let's walk through each part of this implementation. First, parameters should inherit from ParameterBase. This will allow it to automatically determine when its dependencies are met and set the appropriate components of a parent LayerContainer. The second key component is the type annotations on the class itself. Each type annotation represents a dependency of the parameter. Since the original Llama mode has separate query, key, and value dependencies, our fused parameter will declare dependencies for each. Finally, we have the finalize method. This method is automatically called once all dependencies on the layer are met and should return the final parameter.

In this finalize method, we are doing two things: the first is the act of fusing the parameters together through the concatenate method. Note that each of the dependencies can be accessed via self.{name}. The second is calling self.inference_model.transform_qkv_param. A parameter's finalize method always has access to the inference model. In this case we are using that to use a feature provided by DSTransformerBase. This method will automatically shard the parameter for tensor parallelism and then pass it to the linear module implementation to perform additional optimizations or shape transformations, like quantization.

Since many patterns are very common in Transformer models, model_implementations.common_parameters provides implementations for many of the patterns (all compatible with DSTransformerBase) to help accelerate development.

Once all parameters are created, we need to compose them into a layer container. In our simplified Llama model, let's assume there's only QKV and attention output projection matrices. A layer container would appear as the following:

from deepspeed.inference.module_implementations.layer_container_base import LayerContainer

class ExampleContainer(LayerContainer):
    qkvw: UnfusedQKVParameter

    attn_o: AttentionOutputParameter

    PARAM_MAPPING: {
        "self_attn.q_proj.weight": "qkvw.query",
        "self_attn.k_proj.weight": "qkvw.key",
        "self_attn.v_proj.weight": "qkvw.value",
        "self_attn.o_proj.weight": "attn_o.params",
    }

Once again, we have a couple of key components. The first are parameter type annotations. Each annotation corresponds to a parameter that can be used in the model implementation. In the model implementation, I can simply write container.qkvw to access my fused and transformed QKV parameter. The second key component is the PARAM_MAPPING dictionary. This is our explicit mapping of the names of parameters in the source model to a parameter dependency. This mapping dictionary will be used by the policy to automatically populate dependencies.

Once you have written LayerContainers for both the transformer and non-transformer parameters, it's time to work on the model implementation!

Building a Model Implementation that Inherits from DSTransformerBase

By inheriting from DSTransformerBase, most of the implementation work for sharding and transforming parameters will be automatically handled for you. However, there are four key tasks that still need to be completed.

1. Defining the abstract properties based on your model configuration.
2. Configuring embedding and unembedding modules and the forward implementations for them.
3. Configuring the attention configuration and desired KV cache behaviors.
4. Writing the forward implementation for your layer.

Writing a Policy

The InferenceV2Policy is the level of composition. This is the object that will be passed directly to the inference engine and will compose the model implementation and your containers to create an end-to-end solution. There are two main components to be implemented: the first is to create the model that you defined earlier. This is done by implementing the instantiate_model method of the policy. In general, this can just be implemented by calling the constructor for your model and passing the engine config, tensor-parallel communication object, and your custom model config.

The second component is to define how the parameters from the checkpoint will map to each container. From the section on LayerContainers above, you may remember that the LayerContainer can handle the internal routing of a checkpoint parameter to its dependency. In order to find the correct LayerContainer though, we need a second abstraction: the ContainerMap.

A ContainerMap performs this mapping by categorizing checkpoint prefix strings to the type of container they map to. Typically, the easiest way to do this is through iterating over a model checkpoint's state dict or by iterating over the named_parameters of a PyTorch model. There are three types of mappings to define: the transformer mappings, the non-transformer mappings, and the what we'll call the rest. Let's work through an example:

from deepspeed.inference.module_implementations.inference_policy_base import ContainerMap

def build_container_map(self) -> ContainerMap:
    map = ContainerMap()

    transformer_containers = [MyTransformerContainer(self.model) for _ in range(self.model.num_layers)]
    map.set_transformer_params("model.layers", transformer_containers)

    non_transformer_container = MyNonTransformerContainer(self.model)