Créer un backend de modèle compatible pour l'inférence

Les modèles Transformers sont compatibles avec des moteurs d’inférence tels que vLLM et SGLang. Utilisez le même modèle Transformers partout et évitez de réimplémenter un modèle de zéro pour chaque moteur d’inférence. Cela est également utile pour les modèles qu’un moteur ne réimplémente pas nativement, à condition que le modèle prenne en charge l’interface de backend d’attention décrite ci-dessous.

Ce guide vous montre comment implémenter un modèle dans Transformers qui fonctionne comme un moteur de backend pour n'importe quel moteur d'inférence.

Implémentation du modèle

Suivez les consignes de contribution pour les modèles ou les consignes de contribution pour les modèles personnalisés. Le modèle doit contenir un fichier config.json valide dans son répertoire ainsi qu'un champ auto_map valide pointant vers la classe du modèle dans la configuration.

Utilisez la classe AttentionInterface pour les fonctions d'attention personnalisées et optimisées. Cette interface débloque les fonctionnalités de performance de chaque moteur d'inférence.

Utilisez ALL_ATTENTION_FUNCTIONS lors de la définition de la couche d'attention et propagez **kwargs depuis la classe de base MyModel vers les couches d'attention. Définissez _supports_attention_backend sur True dans PreTrainedModel.

Développez le code ci-dessous pour un exemple.

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

Activez le parallélisme tensoriel ou par pipeline optionnel en ajoutant les clés suivantes à PreTrainedConfig.
- base_model_tp_plan active le parallélisme de tenseurs en associant des motifs de noms de couches entièrement qualifiés à des styles de parallélisme de tenseurs. Il ne prend en charge que les stratégies de partitionnement "colwise" et "rowwise".
- base_model_pp_plan active la parallélisation par pipeline en associant les noms des couches enfants directes à des tuples de listes de chaînes de caractères. Le premier élément du tuple contient les noms des arguments d’entrée. Le dernier élément contient les noms des variables de sortie de la couche dans le code de modélisation.
Développez le code ci-dessous pour un exemple.
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"]),
    }
```

Modèles multimodaux

Les modèles multimodaux nécessitent des modifications supplémentaires au-delà des directives de contribution pour les modèles de langage visuels. Ces modifications garantissent que les entrées multimodales sont correctement traitées.

La classe ProcessorMixin doit inclure les attributs self.image_token et self.image_token_ids. Ces jetons de remplacement indiquent les positions des images dans l'entrée. Le même jeton apparaît dans le prompt d'entrée pour les images et dans le code du modèle pour disperser les caractéristiques des images.
La classe ProcessorMixin doit inclure une méthode self._get_num_multimodal_tokens. Cette méthode calcule le nombre de jetons de remplacement nécessaires pour des entrées multimodales de tailles données. Elle renvoie un objet MultiModalData. Les jetons de remplacement situés entre les jetons <image>, tels que les jetons de ligne ou de colonne, ne sont pas comptés comme des jetons de remplacement d’image. Ne comptez que les jetons qui seront remplacés par des caractéristiques d’image plus tard dans le code de modélisation.
La classe ProcessorMixin doit vérifier la valeur de return_mm_token_type_ids et retourner mm_token_type_ids. Cela indique si chaque position correspond à un jeton de texte (0), à un jeton de remplacement d'image (1) ou à un jeton de remplacement de vidéo (2). Les séquences d'identifiants de type de jeton multimodal doivent être contiguës, sans interruption entre les jetons consécutifs. Les jetons spéciaux de début, de fin, de ligne et de colonne doivent être traités comme des jetons de remplacement.

Développez le code ci-dessous pour un exemple.

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)

Ressources

Lisez le billet de blog Intégration du backend Transformers dans vLLM pour plus de détails.
Lisez le billet de blog Intégration du backend Transformers dans SGLang pour plus de détails.

Update on GitHub