새 모델 지원 방법

새 모델 지원 방법 (How to Support New Models)

이 문서는 SGLang에 새로운 언어 모델과 멀티모달 대형 언어 모델(MLLM) 지원을 추가하는 방법을 설명해요. 또한 새 모델을 테스트하고 외부 구현을 등록하는 방법도 다뤄요.

출처: 문서

본문

이 문서는 SGLang에 새로운 언어 모델과 멀티모달 대형 언어 모델(MLLM) 지원을 추가하는 방법을 설명해요. 또한 새 모델을 테스트하고 외부 구현을 등록하는 방법도 다뤄요.

새 언어 모델 지원 방법 (How to Support a New Language Model)

SGLang에서 새 모델을 지원하려면 SGLang Models Directory 아래에 파일 하나만 추가하면 돼요. 기존 모델 구현에서 배우고 모델용 새 파일을 만들 수 있어요. 대부분의 모델에서 시작할 유사한 모델을 찾을 수 있을 거예요(예: Llama에서 시작). 또한 vLLM에서 SGLang으로 모델 포팅하는 방법도 참고하세요.

새 멀티모달 대형 언어 모델 지원 방법 (How to Support a New Multimodal Large Language Model)

SGLang에서 새 멀티모달 대형 언어 모델(MLLM)을 지원하려면 표준 LLM 지원 외에 몇 가지 핵심 컴포넌트가 필요해요:

  1. 새 모델을 멀티모달로 등록: model_config.py의 is_multimodal_model을 확장해 모델에 대해 True를 반환하도록 해요.

  2. 새 chat-template 등록: 기본 chat-template이 이미지를 입력으로 받지 못할 때만: conversation.py에 새 chat template과 해당 매칭 함수를 등록해요.

  3. 멀티모달 데이터 프로세서: BaseMultimodalProcessor에서 상속하는 새 Processor 클래스를 정의하고 이 프로세서를 모델의 전용 프로세서로 등록해요. 자세한 내용은 multimodal_processor.py를 참고해요.

  4. 멀티모달 토큰 처리: 새 모델용 pad_input_ids 함수를 구현해요. 이 함수에서 프롬프트의 멀티모달 토큰을 (필요하면) 확장하고 멀티모달 데이터 해시로 패딩해서, SGLang이 RadixAttention으로 서로 다른 멀티모달 데이터를 인식할 수 있게 해요.

  5. 이미지 피처 추출 처리: 새 모델용 get_image_feature 함수를 구현해요. 이 함수는 원시 이미지 데이터에서 이미지 피처를 추출해 언어 모델에서 사용하는 임베딩으로 변환해요.

  6. 비전 어텐션 적응: ViT의 멀티 헤드 Attention을 SGLang의 VisionAttention에 적응시켜요.

Qwen2VL이나 다른 mllm 구현을 참고할 수 있어요. 이 모델들은 멀티모달 및 텍스트 입력을 모두 올바르게 처리하는 방법을 보여줘요.

테스트와 디버깅 (Testing and Debugging)

모든 테스트와 벤치마크 결과를 PR 설명에 기록해 주세요.

대화형 디버깅 (Interactive Debugging)

대화형 디버깅에서는 Hugging Face/Transformers와 SGLang의 출력을 비교해요. 다음 두 명령은 동일한 텍스트 출력과 매우 유사한 prefill logits를 제공해야 해요:

  • 참조 출력 얻기: 
    python3 scripts/playground/reference_hf.py --model-path [new model] --model-type {text,vlm}
  • SGLang 출력 얻기: 
    python3 -m sglang.benchmark.one_batch --correct --model [new model]

테스트 스위트에 모델 추가 (Add the Model to the Test Suite)

새 모델이 잘 유지되도록 하려면 test_generation_models.py 파일의 ALL_OTHER_MODELS 목록에 추가하고, 로컬 머신에서 새 모델을 테스트하며 시범 벤치마크(GSM8K, MMLU, MMMU, MMMU-Pro 등)의 결과를 PR에 보고해요. VLM의 경우 test_vision_openai_server_{x}.py(예: test_vision_openai_server_a.py)에도 테스트를 포함해요.

로컬 머신에서 새 모델을 테스트하는 예시 명령:

ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels.test_others

벤치마크 (Benchmark)

  • (필수) MMMU: MMMU 벤치마크 README.md를 따라 SGLang vs HF Transformer 정확도 비교를 얻어요. SGLang 실행의 정확도 점수는 HF Transformer 실행보다 훨씬 낮아선 안 돼요. 비슷하게 https://docs.sglang.io/developer_guide/benchmark_and_profiling.html를 따라 성능 비교를 얻어요: TTFT와 처리량이 기준선(예: HF Transformer)을 충족하거나 초과해야 해요.
  • (선택) 기타 평가: 다른 평가를 실행했다면 PR 설명에 결과를 기록해 주세요.

vLLM에서 SGLang으로 모델 포팅 (Port a Model from vLLM to SGLang)

vLLM Models Directory는 vLLM이 많은 모델을 다루므로 귀중한 자료예요. SGLang은 vLLM의 인터페이스와 일부 레이어를 재사용하므로 vLLM에서 SGLang으로 모델을 포팅하기가 더 쉬워요.

vLLM에서 SGLang으로 모델을 포팅하려면:

  • 이 두 파일을 비교해서 참고해요:
  • 주요 차이점은 다음과 같아요:
    • vLLM의 Attention을 RadixAttention으로 교체 (layer_id를 RadixAttention에 전달해야 해요).
    • vLLM의 LogitsProcessor를 SGLang의 LogitsProcessor로 교체.
    • ViT의 멀티 헤드 Attention을 SGLang의 VisionAttention으로 교체.
    • 다른 vLLM 레이어(예: RMSNorm, SiluAndMul)를 SGLang 레이어로 교체.
    • Sample 제거.
    • forward() 함수 변경 및 forward_batch() 메서드 추가.
    • 마지막에 EntryClass 추가.
    • 새 구현이 SGLang 컴포넌트만 사용하고 vLLM 컴포넌트에 의존하지 않도록 보장.

참고: 새 모델을 supported models 문서의 지원 모델 목록에 추가해 주세요.

외부 모델 구현 등록 (Registering an External Model Implementation)

위 방법 외에도 서버를 시작하기 전에 ModelRegistry로 새 모델을 등록할 수 있어요. 이렇게 하면 소스 코드를 수정하지 않고 모델을 통합할 수 있어요.

예를 들면:

from sglang.srt.models.registry import ModelRegistry
from sglang.srt.entrypoints.http_server import launch_server

# For a single model, add it to the registry:
ModelRegistry.models[model_name] = model_class

# For multiple models, you can imitate the import_model_classes() function:
from functools import lru_cache

@lru_cache()
def import_new_model_classes():
    model_arch_name_to_cls = {}
    # Populate model_arch_name_to_cls with your new model classes.
    ...
    return model_arch_name_to_cls

ModelRegistry.models.update(import_new_model_classes())

# Launch the server with your server arguments:
launch_server(server_args)

예시: Llama 래퍼 모델 구현 및 서빙 (Example: Implementing and Serving a Llama Wrapper Model)

아래는 SGLang에서 새 모델을 엔드투엔드로 구현하고 Offline Engine으로 실행하는 초보자용 단계별 안내예요.

모델 구현 (Implementing Our Model)

간단히 하기 위해 이 새 모델은 Llama 3.1-8B-Instruct를 감싸는 간단한 래퍼이며, 각 forward 호출에서 개별 로짓의 제곱근을 취해 출력 로짓을 바이어스하는 것을 목표로 해요.

llama_wrapper.py라는 파일에 모델을 정의해 보아요. 첫 단계는 SGLang의 내부 백엔드인 SRT에서 필요한 라이브러리를 가져오는 거예요.

# In the file `llama_wrapper.py`

import torch
from transformers import LlamaConfig
from typing import Optional
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors

from sglang.srt.models.llama import LlamaForCausalLM

다음으로 모델을 위한 새 class를 선언하고 LlamaForCausalLM에서 상속하게 해요. 이렇게 하면 LlamaAttention, LlamaMLP 같은 LlamaForCausalLM의 사전 정의된 모듈과 레이어에 접근할 수 있어요. 거의 모든 모델 구현은 __init__ 메서드의 인자로 config와 quant_config를 받는다는 점을 참고하세요. config와 quant_config는 model_loader/loader.py를 통해 전달돼요. LlamaForCausalLM에서 상속했으므로 파라미터를 직접 그 생성자에 전달할 수 있고, 생성자가 멤버 변수를 설정해 줘요.

class LlamaWrapper(LlamaForCausalLM):
    def __init__(
        self,
        config: LlamaConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__(config=config, quant_config=quant_config, prefix=prefix)

이제 추론 시 호출될 forward 메서드를 정의해요. forward의 시그니처는 거의 모든 모델에서 본질적으로 동일하다는 점을 참고하세요. 참고는 models 디렉터리에 정의된 다른 모델들을 보세요. SGLang 런타임 내부에서 forward가 정확히 어디서 호출되는지 보려면 ModelRunner 클래스의 forward_decode와 forward_extend를 보세요.

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
        input_embeds: Optional[torch.Tensor] = None,
        get_embedding: bool = False,
    ) -> LogitsProcessorOutput:

이제 self.model( LlamaForCausalLM이 __init__에서 정의하는 멤버 변수)의 __call__ 메서드를 호출하며, 이는 결국 LlamaForCausalLM의 forward 메서드를 호출해요. 그 후 hidden_states를 모델의 LogitsProcessor(LlamaForCausalLM에서 정의)에 넣어요.

        hidden_states = self.model(
            input_ids,
            positions,
            forward_batch,
            input_embeds,
            pp_proxy_tensors=pp_proxy_tensors,
        )

        res: LogitsProcessorOutput = self.logits_processor(
            input_ids,
            hidden_states,
            self.lm_head,
            forward_batch,
        )

다음 토큰의 로짓을 받은 뒤 마침내 바이어싱 단계를 수행할 수 있어요.

        orig_logits = res.next_token_logits
        res.next_token_logits = torch.where(
            orig_logits > 0,
            orig_logits.sqrt(),
            orig_logits
        )

        return res

이제 LlamaWrapper 모델이 생성되어 서빙 준비가 됐어요!

SGLang의 Offline Engine으로 모델 서빙 (Serving Our Model Via SGLang's Offline Engine)

이 안내의 다음 단계는 새 모델을 오프라인으로 호스팅해서 HTTP 서버 없이 로컬에서 서빙할 수 있게 하는 거예요.

먼저 run.py라는 새 파일을 만들어요. 이제 SGLang의 ModelRegistry가 우리 모델을 찾을 수 있도록 해야 해요. 이를 위해 먼저 Huggingface에서 모델의 구성과 가중치를 내려받아요.

# In the file `run.py`

import asyncio
from functools import lru_cache
from huggingface_hub import snapshot_download
from llama_wrapper import LlamaWrapper # Make sure to import our new model!
import sglang as sgl
from sglang.srt.models.registry import ModelRegistry

# Make sure to request access to this model on Huggingface, then export your
# `HF_TOKEN` to download the model snapshot
llama_dir = snapshot_download(
    repo_id="meta-llama/Llama-3.1-8B-Instruct",
    local_dir="./llama_ckpt",
)

이제 모델이 디스크에 있으니 ./llama_ckpt/config.json의 architectures 필드를 LlamaWrapper로 변경해 LlamaWrapper를 가리키게 해요. 그렇게 하면 모델 체크포인트 경로를 SGLang에 전달할 때, SGLang은 모델로 "LlamaForCausalLM" 대신 "LlamaWrapper"를 사용하려는 것임을 알게 돼요.

{
  "architectures": [
   #  "LlamaForCausalLM"
    "LlamaWrapper"
  ],
  ...
}

하지만 LlamaWrapper 클래스를 "LlamaWrapper" 레지스트리 키워드로 연결하지 않으면 SGLang이 모델을 찾을 수 없어요. 따라서 LlamaWrapper를 등록하려면 위 "Registering an External Model Implementation" 섹션의 단계를 따라야 해요.

@lru_cache()
def import_new_model_classes():
    model_arch_name_to_cls = {"LlamaWrapper": LlamaWrapper}
    return model_arch_name_to_cls

ModelRegistry.models.update(import_new_model_classes())

마지막으로 Engine을 만들 때 로컬 모델 디렉터리 경로를 전달하면 돼요. 그러면 LlamaWrapper가 서빙 준비가 돼요. 이 안내에서는 SGLang Engine의 비스트리밍 비동기 생성 엔드포인트를 사용할게요.

def main():
    llm = sgl.Engine(model_path="./llama_ckpt")
    sampling_params = {"temperature": 0.2, "top_k": 5}
    prompts = [
        "Write a short, neutral self-introduction for a fictional character. Hello, my name is",
        "Provide a concise factual statement about France’s capital city. The capital of France is",
        "Explain possible future trends in artificial intelligence. The future of AI is",
    ]

    asyncio.run(run_llm(llm, sampling_params, prompts))

    llm.shutdown()

async def run_llm(
    llm,
    sampling_params,
    prompts,
) -> None:
    outputs = await llm.async_generate(prompts, sampling_params)

    for prompt, output in zip(prompts, outputs):
        print(f"\nPrompt: {prompt}")
        print(f"Generated text: {output['text']}")

if __name__ == "__main__":
    main()

이제 python run.py를 호출하면 새로 만든 모델의 출력을 얻을 수 있어요!

표준 CLI로 외부 모델 서빙 (Serving External Models via the Standard CLI)

이전 섹션들은 ModelRegistry로 모델을 프로그래밍 방식으로 등록하고 Offline Engine으로 서빙하는 방법을 보여줘요. vLLM 모델 플러그인과 유사하게, SGLang 소스 코드를 수정하지 않고 표준 python -m sglang.launch_server CLI를 계속 사용할 수 있는 대안이 있어요: SGLANG_EXTERNAL_MODEL_PACKAGE 환경 변수로 모델을 등록하면 돼요.

EntryClass 변수

SGLang이 모델 패키지를 스캔할 때 Python 파일의 모듈 레벨에서 EntryClass 변수를 찾아요. 모델 레지스트리가 파일을 가져와 EntryClass를 확인하고 그에 할당된 클래스를 등록해요. HuggingFace 기반 모델을 사용한다면 이 클래스 이름은 모델 config.json의 "architectures" 필드와 일치해야 해요.

예를 들어 Llama 래퍼를 구현한다면 모델 파일 끝에 이 줄을 추가해요:

# This is what "Add EntryClass at the end" means
EntryClass = LlamaWrapper

예시: 텍스트 전용 모델 (Example: Text-Only Model)

이전 섹션과 동일한 Llama 래퍼를 사용해 CLI로 패키징하고 서빙하는 방법이에요.

  1. 프로젝트 생성
sglang_custom_project/
|----setup.py
|----custom_llm/
     |----__init__.py
     |----llama_wrapper.py

setup.py 작성:

# sglang_custom_project/setup.py

from setuptools import setup, find_packages
setup(
    name="sglang-custom-plugins",
    version="0.1",
    packages=find_packages(),
)
  1. 모델 코드 작성

llama_wrapper.py 안에 모델을 작성하고 EntryClass를 포함해요:

# sglang_custom_project/custom_llm/llama_wrapper.py

import torch
from typing import Optional
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.models.llama import LlamaForCausalLM

class LlamaWrapper(LlamaForCausalLM):
    def __init__(self, config, quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = "") -> None:
        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
    @torch.no_grad()
    def forward(self, input_ids, positions, forward_batch,
                pp_proxy_tensors=None, input_embeds=None, get_embedding=False):
        hidden_states = self.model(
            input_ids, positions, forward_batch, input_embeds,
            pp_proxy_tensors=pp_proxy_tensors,
        )
        res: LogitsProcessorOutput = self.logits_processor(
            input_ids, hidden_states, self.lm_head, forward_batch,
        )

        orig = res.next_token_logits
        res.next_token_logits = torch.where(orig > 0, orig.sqrt(), orig)
        return res

# Don't forget to add EntryClass
EntryClass = LlamaWrapper
  1. 패키지 설치

sglang_custom_project 디렉터리 안에서 이것을 실행해 활성 Python 환경에 코드를 설치해요:

pip install -e .
  1. config.json 업데이트

HuggingFace 모델 체크포인트 디렉터리 아래의 config.json에서 architectures 필드가 클래스 이름과 일치하도록 업데이트해요:

{
  "architectures": ["LlamaWrapper"],
  ...
}
  1. 서버 실행

CLI를 실행하기 전에 환경 변수를 설정해요:

export SGLANG_EXTERNAL_MODEL_PACKAGE=custom_llm
python -m sglang.launch_server \
    --model-path /path/to/Llama-3.1-8B-Instruct \
    --port 8000

SGLANG_EXTERNAL_MODEL_PACKAGE는 모델 관련 코드를 포함하는 상위 폴더 이름이어야 해요. 이 예시에서는 custom_llm이어야 해요.

예시: 멀티모달 모델 (Example: Multimodal Model)

멀티모달 모델로 작업한다면 SGLANG_EXTERNAL_MODEL_PACKAGE만 설정하는 것으로는 충분하지 않아요. SGLang은 이미지/비디오 처리 파이프라인을 활성화하려면 아키텍처를 멀티모달로 인식해야 하고, 커스텀 프로세서도 필요해요.

추가 환경 변수 두 개로 처리할 수 있어요:

  • SGLANG_EXTERNAL_MM_MODEL_ARCH: 아키텍처 이름을 SGLang의 내부 멀티모달 모델 목록에 추가해요.
  • SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE: 커스텀 프로세서 클래스를 어디서 찾을지 SGLang에 알려줘요.

예를 들어 로짓의 제곱근을 취하는 Qwen2-VL-Instruct 기반 커스텀 모델을 만들어 보아요.

프로젝트 생성:

sglang_custom_project_vl/
|----setup.py
|----custom_vlm/
     |----__init__.py
     |----qwenvl_wrapper.py

setup.py 작성:

# sglang_custom_project_vl/setup.py

from setuptools import setup, find_packages
setup(
    name="sglang-custom-plugins-vl",
    version="0.1",
    packages=find_packages(),
)

qwenvl_wrapper.py에 모델 작성:

# sglang_custom_project_vl/custom_vlm/qwenvl_wrapper.py
import torch
from sglang.srt.models.qwen2_vl import Qwen2VLForConditionalGeneration
from sglang.srt.multimodal.processors.qwen_vl import QwenVLImageProcessor

class CustomQwen2VL(Qwen2VLForConditionalGeneration):
    def forward(self, input_ids, positions, forward_batch,
                input_embeds=None, get_embedding=False):
        res = super().forward(
            input_ids, positions, forward_batch,
            input_embeds=input_embeds, get_embedding=get_embedding
        )
        if not get_embedding:
            orig = res.next_token_logits
            res.next_token_logits = torch.where(orig > 0, orig.sqrt(), orig)
        return res

class CustomQwen2VLProcessor(QwenVLImageProcessor):
    models = [CustomQwen2VL]

    def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
        super().__init__(hf_config, server_args, _processor, *args, **kwargs)

EntryClass = CustomQwen2VL

Note: 프로세서를 특정 모델 클래스와 연결하는 한 커스텀 프로세서용 별도 EntryClass는 필요 없어요.

패키지 설치, config.json 업데이트, 실행:

pip install -e .

{
  "architectures": ["CustomQwen2VL"],
  ...
}

export SGLANG_EXTERNAL_MODEL_PACKAGE=custom_vlm
export SGLANG_EXTERNAL_MM_MODEL_ARCH=CustomQwen2VL
export SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE=custom_vlm

python -m sglang.launch_server \
    --model-path /path/to/Qwen2-VL-2B-Instruct \
    --port 8000 \
    --enable-multimodal

문서화 (Documentation)

generative_models.md 또는 multimodal_language_models.md의 지원 모델 표에 추가해요.

이 지침을 따르면 SGLang에서 새 언어 모델과 멀티모달 대형 언어 모델 지원을 추가하고, 철저히 테스트되며 시스템에 쉽게 통합되도록 보장할 수 있어요.

더 알아보기