Criando um backend de modelo compatível para inferência

Os modelos do Transformers são compatíveis com motores de inferência como vLLM e SGLang. Use o mesmo modelo do Transformers em qualquer lugar e evite reimplementar um modelo do zero para cada motor de inferência. Isso também ajuda com modelos que um motor não reimplementa nativamente, desde que o modelo suporte a interface de backend de atenção descrita abaixo.

Este guia mostra como implementar um modelo em Transformers que funciona como um backend para qualquer mecanismo de inferência.

Implementação do modelo

Siga as diretrizes de contribuição de modelos ou as diretrizes de contribuição de modelos personalizados. O modelo deve ter um config.json válido em seu diretório e um campo auto_map válido apontando para a classe do modelo na configuração.

Use a classe AttentionInterface para funções de atenção personalizadas e otimizadas. Essa interface libera os recursos de desempenho de cada mecanismo de inferência.

Use ALL_ATTENTION_FUNCTIONS ao definir a camada de atenção e propague **kwargs da classe base MyModel para as camadas de atenção. Defina _supports_attention_backend como True em PreTrainedModel.

Expanda o código abaixo para um exemplo.

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

Ative a paralelização opcional de tensor ou de pipeline adicionando as seguintes chaves a PreTrainedConfig.

base_model_tp_plan habilita o tensor parallelism mapeando padrões de nomes de camadas totalmente qualificados para estilos de paralelismo de tensor. Suporta apenas as estratégias de particionamento "colwise" e "rowwise".
base_model_pp_plan habilita a paralelização de pipeline mapeando os nomes das camadas filhas diretas para tuplas de listas de strings. O primeiro elemento da tupla contém os nomes dos argumentos de entrada. O último elemento contém os nomes das variáveis das saídas da camada no código de modelagem.

Expanda o código abaixo para um exemplo.

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 multimodais

Modelos multimodais exigem alterações adicionais além das diretrizes de contribuição para modelos de linguagem visual. Essas alterações garantem que as entradas multimodais sejam processadas corretamente.

A classe ProcessorMixin deve incluir os atributos self.image_token e self.image_token_ids. Esses tokens de espaço reservado indicam as posições das imagens na entrada. O mesmo token aparece no prompt de entrada para as imagens e no código do modelo para dispersar os recursos de imagem.
A classe ProcessorMixin deve incluir um método self._get_num_multimodal_tokens. Este método calcula o número de tokens de espaço reservado necessários para entradas multimodais com tamanhos específicos. Ele retorna um objeto MultiModalData. Os espaços reservados entre os tokens <image>, como tokens de linha ou coluna, não são contados como espaços reservados de imagem. Conte apenas os tokens que serão substituídos por recursos de imagem mais adiante no código de modelagem.
A classe ProcessorMixin deve verificar o valor de return_mm_token_type_ids e retornar mm_token_type_ids. Isso indica se cada posição é um token de texto (0), um token placeholder de imagem (1) ou um token placeholder de vídeo (2). As sequências de ids de tipo de token multimodal devem ser contíguas, sem interrupções entre tokens consecutivos. Trate os tokens especiais de início, fim, linha e coluna como placeholders.

Expanda o código abaixo para um exemplo.

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

Leia o post no blog Integração do backend Transformers no vLLM para mais detalhes.
Leia o post no blog Integração do backend Transformers no SGLang para mais detalhes.

Update on GitHub