Construir un backend de modelo compatible para inferencia

Los modelos de Transformers son compatibles con motores de inferencia como vLLM y SGLang. Utilice el mismo modelo de Transformers en cualquier lugar y evite reimplementar un modelo desde cero para cada motor de inferencia. Esto también resulta útil para los modelos que un motor no reimplementa de forma nativa, siempre que el modelo admita la interfaz de backend de atención descrita a continuación.

Esta guía le muestra cómo implementar un modelo en Transformers que funcione como backend para cualquier motor de inferencia.

Implementación del modelo

Siga las directrices de contribución de modelos o las directrices de contribución de modelos personalizados. El modelo debe tener un config.json válido en su directorio y un campo auto_map válido que apunte a la clase del modelo en la configuración.

Utilice la clase AttentionInterface para funciones de atención personalizadas y optimizadas. Esta interfaz habilita las características de rendimiento de cada motor de inferencia.

Use ALL_ATTENTION_FUNCTIONS al definir la capa de atención y propague **kwargs desde la clase base MyModel a las capas de atención. Establezca _supports_attention_backend en True en PreTrainedModel.

Expanda el siguiente código para un ejemplo.

modeling_my_model.py

from transformers import PreTrainedModel
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from torch import nn

class MyAttention(nn.Module):

    def forward(self, hidden_states, **kwargs):
        ...
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            **kwargs,
        )
        ...

class MyModel(PreTrainedModel):
    _supports_attention_backend = True

Habilite el paralelismo de tensor o de pipeline opcional añadiendo las siguientes claves a PreTrainedConfig.
- base_model_tp_plan habilita el paralelismo de tensores mediante el mapeo de patrones de nombres de capas totalmente cualificados a estilos de paralelismo de tensores. Admite únicamente las estrategias de partición "colwise" y "rowwise".
- base_model_pp_plan habilita la paralelización por pipeline mediante el mapeo de los nombres de las capas hijas directas a tuplas de listas de cadenas de caracteres. El primer elemento de la tupla contiene los nombres de los argumentos de entrada. El último elemento contiene los nombres de las variables de las salidas de la capa en el código de modelado.
Expanda el siguiente código para un ejemplo.
configuration_my_model.py
```
from transformers import PreTrainedConfig

class MyConfig(PreTrainedConfig):
    base_model_tp_plan = {
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }
```

Modelos multimodales

Los modelos multimodales requieren cambios adicionales además de las directrices de contribución para modelos de lenguaje visual. Estos cambios garantizan que las entradas multimodales se procesen correctamente.

La clase ProcessorMixin debe incluir los atributos self.image_token y self.image_token_ids. Estos tokens de marcador de posición indican las posiciones de las imágenes en la entrada. El mismo token aparece en el prompt de entrada para las imágenes y en el código del modelo para dispersar las características de las imágenes.
La clase ProcessorMixin debe incluir un método self._get_num_multimodal_tokens. Este método calcula el número de tokens de marcador de posición necesarios para entradas multimodales con tamaños dados. Devuelve un objeto MultiModalData. Los marcadores de posición entre los tokens <image>, como los tokens de fila o columna, no se cuentan como marcadores de posición de imagen. Cuenta solo los tokens que se reemplazan por características de imagen más adelante en el código de modelado.
La clase ProcessorMixin debe comprobar el valor de return_mm_token_type_ids y devolver mm_token_type_ids. Esto indica si cada posición es un token de texto (0), un token de marcador de posición de imagen (1) o un token de marcador de posición de vídeo (2). Las secuencias de id de tipo de token multimodal deben ser contiguas, sin interrupciones entre tokens consecutivos. Trate los tokens especiales de inicio, fin, fila y columna como marcadores de posición.

Expanda el siguiente código para un ejemplo.

modeling_my_multimodal_model.py

class MyMultimodalProcessor(ProcessorMixin):

    def __call__(self, images=None, text=None, **kwargs):
        if return_mm_token_type_ids:
            mm_token_type_ids = np.zeros_like(input_ids)
            mm_token_type_ids[input_ids == self.image_token_id] = 1
            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
        Args:
            image_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (height, width) per each image.
        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """
        vision_data = {}
        if image_sizes is not None:
            num_image_tokens = [256] * len(image_sizes) # 256 placeholder tokens for each image always
            num_image_patches = [1] * len(image_sizes) # no patching, thus each image is processed as a single base image
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
        return MultiModalData(**vision_data)

Recursos

Lea la publicación de blog Integración del backend de Transformers en vLLM para obtener más detalles.
Lea la publicación de blog Integración del backend de Transformers en SGLang para obtener más detalles.

Update on GitHub