추론용 호환 모델 백엔드 만들기

추론용 호환 모델 백엔드 만들기

Transformers 모델은 vLLM, SGLang 같은 추론 엔진의 백엔드로 동작할 수 있어요. 추론 엔진의 백엔드가 될 수 있는 모델을 Transformers에서 구현하는 방법을 살펴볼게요.

출처: 문서

본문

Transformers 모델은 vLLM과 SGLang 같은 추론 엔진과 호환돼요. 같은 Transformers 모델을 어디서나 사용해서, 추론 엔진마다 모델을 처음부터 다시 구현하지 않아도 돼요. 엔진이 네이티브로 재구현하지 않는 모델에서도, 아래 설명하는 attention 백엔드 인터페이스를 지원하기만 하면 도움이 돼요.

이 가이드는 어떤 추론 엔진의 백엔드로도 동작하는 모델을 Transformers에서 구현하는 방법을 보여줘요.

모델 구현 (Model implementation)

  1. 모델 기여 가이드라인 또는 커스텀 모델 기여 가이드라인을 따라요. 모델의 디렉터리에 유효한 config.json이 있어야 하고, config에 모델 클래스를 가리키는 유효한 auto_map 필드가 있어야 해요.

  2. 커스텀·최적화 attention 함수에는 AttentionInterface 클래스를 사용해요. 이 인터페이스가 각 추론 엔진의 성능 기능을 열어줘요.

    attention 레이어를 정의할 때는 ALL_ATTENTION_FUNCTIONS를 사용하고, 기본 MyModel 클래스에서 attention 레이어로 **kwargs를 전파해요. PreTrainedModel에서 _supports_attention_backend를 True로 설정해요.

    아래 코드를 펼치면 예시가 나와요.

    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

  3. PreTrainedConfig에 다음 키를 추가해 선택적으로 텐서 또는 파이프라인 병렬화를 활성화해요.

    • base_model_tp_plan은 완전히 정규화된 레이어 이름 패턴을 텐서 병렬화 스타일로 매핑해 텐서 병렬화를 활성화해요. "colwise"와 "rowwise" 분할 전략만 지원해요.
    • base_model_pp_plan은 직접 자식 레이어 이름을 문자열 리스트의 튜플에 매핑해 파이프라인 병렬화를 활성화해요. 튜플의 첫 번째 요소는 입력 인자 이름을 담고, 마지막 요소는 모델링 코드에서 레이어 출력의 변수 이름을 담아요.

    아래 코드를 펼치면 예시가 나와요.

    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"]),
    }

멀티모달 모델 (Multimodal models)

멀티모달 모델은 비전 언어 모델 기여 체크리스트 이상의 추가 변경이 필요해요. 이러한 변경은 멀티모달 입력이 올바르게 처리되도록 보장해요.

  1. ProcessorMixin 클래스에 self.image_token과 self.image_token_ids 속성이 포함돼야 해요. 이 플레이스홀더 토큰들은 입력에서 이미지 위치를 나타내요. 같은 토큰이 이미지 입력 프롬프트와 모델 코드(이미지 피처를 흩뿌리는) 양쪽에 나타나요.

  2. ProcessorMixin 클래스에 self._get_num_multimodal_tokens 메서드가 포함돼야 해요. 이 메서드는 주어진 크기의 멀티모달 입력에 필요한 플레이스홀더 토큰 수를 계산해요. MultiModalData 객체를 반환해요. <image> 토큰 사이의 플레이스홀더(예: 행, 열 토큰)는 이미지 플레이스홀더로 세지 않아요. 나중에 모델링 코드에서 이미지 피처로 대체되는 토큰만 세요.

  3. ProcessorMixin 클래스가 return_mm_token_type_ids 값을 확인하고 mm_token_type_ids를 반환해야 해요. 이는 각 위치가 텍스트 토큰(0), 이미지 플레이스홀더 토큰(1), 또는 비디오 플레이스홀더 토큰(2)인지를 나타내요. 멀티모달 토큰 타입 id 시퀀스는 연속된 토큰 사이에 끊김이 없어야 해요. 시작, 끝, 행, 열 토큰 같은 특수 토큰은 플레이스홀더로 취급해요.

아래 코드를 펼치면 예시가 나와요.

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)

리소스 (Resources)

더 알아보기 (Learn more)