Triton Inference Server와 제약 디코딩(Constrained Decoding)

이 튜토리얼은 제약 디코딩(Constrained Decoding) 을 다뤄요. 대규모 언어 모델(LLM)이 생성하는 출력을 엄격한 포맷 요구 사항에 맞추도록 강제하는 중요한 기법인데, 파인튜닝만으로는 맞추기 어렵거나 비용이 큰 요구 사항이죠.

제약 디코딩 소개

제약 디코딩은 자연어 처리와 다양한 AI 애플리케이션에서 모델의 출력을 안내·제어하는 강력한 기법이에요. 특정 제약을 부과해 생성 출력이 길이, 포맷, 내용 제한 같은 미리 정의된 기준을 지키도록 보장합니다. 유효한 코드 조각, 구조화된 데이터, 문법적으로 올바른 문장을 만들어야 하는 것처럼 규칙 준수가 필수적인 상황에서 핵심적인 능력이죠.

최근에는 일부 모델이 애초에 이런 제약을 내재하도록 파인튜닝되기도 해요. 생성 과정에서 제약을 자연스럽게 통합해 광범위한 후처리가 필요 없어지고, 정확성과 신뢰성이 중요한 자동 콘텐츠 생성·데이터 검증·실시간 번역 같은 작업에서 가치가 큽니다.

이 튜토리얼은 이미 JSON 구조화 출력을 지원하는 Hermes-2-Pro-Llama-3-8B를 기반으로 해요. Hermes-2-Pro-Llama-3-8B를 Triton Inference Server와 TensorRT-LLM 백엔드로 배포하는 자세한 과정은 이 튜토리얼에서 볼 수 있습니다. 이런 경우 출력의 구조·품질은 프롬프트 엔지니어링으로 제어할 수 있고, 그 경로는 '프롬프트 엔지니어링을 통한 구조화 생성' 절에서 다룹니다.

제약 디코딩에 파인튜닝되지 않은 모델이거나 출력을 더 정밀하게 제어하고 싶다면 LM Format Enforcer와 Outlines 같은 전용 라이브러리가 강력한 해법이 돼요. 이 라이브러리들은 모델 출력에 특정 제약을 강제해, 원하는 기준에 정확히 맞게 생성 과정을 조정하게 해 줍니다. 이 튜토리얼에서는 워크플로에 LM Format Enforcer와 Outlines를 쓰는 법을 보여줄게요.

사전 준비: Hermes-2-Pro-Llama-3-8B

진행 전에 Hermes-2-Pro-Llama-3-8B 모델을 이 단계를 따라 Triton Inference Server와 TensorRT-LLM 백엔드로 성공적으로 배포했는지 확인하세요.

프롬프트 엔지니어링을 통한 구조화 생성

먼저 Triton SDK 컨테이너를 시작해요.

# SDK 컨테이너 예시
docker run --rm -it --net host --shm-size=2g \
    --ulimit memlock=-1 --ulimit stack=67108864 --gpus all \
    -v /path/to/tutorials:/tutorials \
    -v /path/to/Hermes-2-Pro-Llama-3-8B/repo:/Hermes-2-Pro-Llama-3-8B \
    nvcr.io/nvidia/tritonserver:<xx.yy>-py3-sdk

제공되는 클라이언트 스크립트는 pydantic 라이브러리를 쓰는데, SDK 컨테이너에는 포함되어 있지 않아요. 진행 전에 설치하세요.

pip install pydantic

예시 1

파인튜닝된 모델은 간단히 시스템 프롬프트만으로 JSON 모드를 켤 수 있어요.

You are a helpful assistant that answers in JSON.

전체 prompt 구성 로직은 client.py를 참고하세요.

python3 /tutorials/AI_Agents_Guide/Constrained_Decoding/artifacts/client.py --prompt "Give me information about Harry Potter and the Order of Phoenix" -o 200 --use-system-prompt

다음과 같은 응답을 기대할 수 있어요.

...
assistant
{
  "title": "Harry Potter and the Order of Phoenix",
  "book_number": 5,
  "author": "J.K. Rowling",
  "series": "Harry Potter",
  "publication_date": "June 21, 2003",
  "page_count": 766,
  "publisher": "Arthur A. Levine Books",
  "genre": [
    "Fantasy",
    "Adventure",
    "Young Adult"
  ],
  "awards": [
    {
      "award_name": "British Book Award",
      "category": "Children's Book of the Year",
      "year": 2004
    }
  ],
  "plot_summary": "Harry Potter and the Order of Phoenix is the fifth book in the Harry Potter series. In this installment, Harry returns to Hogwarts School of Witchcraft and Wizardry for his fifth year. The Ministry of Magic is in denial about the return of Lord Voldemort, and Harry finds himself battling against the"

예시 2

선택적으로 출력을 특정 스키마로 제한할 수도 있어요. 예를 들어 client.py에서 pydantic 라이브러리로 다음 답변 형식을 정의합니다.

from pydantic import BaseModel

class AnswerFormat(BaseModel):
    title: str
    year: int
    director: str
    producer: str
    plot: str

...

prompt += "Here's the json schema you must adhere to:\n<schema>\n{schema}\n</schema>".format(
                schema=AnswerFormat.model_json_schema())

실행해 볼게요.

python3 /tutorials/AI_Agents_Guide/Constrained_Decoding/artifacts/client.py --prompt "Give me information about Harry Potter and the Order of Phoenix" -o 200 --use-system-prompt --use-schema

다음과 같은 응답을 기대할 수 있어요.

 ...
assistant
{
  "title": "Harry Potter and the Order of Phoenix",
  "year": 2007,
  "director": "David Yates",
  "producer": "David Heyman",
  "plot": "Harry Potter and his friends must protect Hogwarts from a threat when the Ministry of Magic is taken over by Lord Voldemort's followers."
}

외부 라이브러리로 출력 포맷 강제하기

이 절에서는 제약 디코딩에 파인튜닝되지 않은 LLM에 제약을 부과하는 법을 보여줄게요. LM Format Enforcer와 Outlines가 강력한 해법입니다.

두 라이브러리의 참조 구현은 출력 포맷인 AnswerFormat을 정의하는 utils.py 스크립트에 담겨 있어요.

class WandFormat(BaseModel):
        wood: str
        core: str
        length: float

class AnswerFormat(BaseModel):
        name: str
        house: str
        blood_status: str
        occupation: str
        alive: str
        wand: WandFormat

사전 준비: 공통 설정

이 단계를 따라 Hermes-2-Pro-Llama-3-8B 모델을 Triton Inference Server와 TensorRT-LLM 백엔드로 성공적으로 배포했는지 확인하세요.

중요: 도커 컨테이너를 시작할 때 tutorials 폴더를 /tutorials에 마운트했는지 확인하세요.

설정이 끝나면 /opt/tritonserver/inflight_batcher_llm 폴더가 있어야 하고, 추론 요청(예시 1·2의 것)을 몇 번 시도해 볼 수 있어야 해요.

모델 파일을 조정할 것이므로, 실행 중인 서버가 있다면 멈출 수 있어요.

pkill tritonserver

Logits 후처리기(Post-Processor)

두 라이브러리 모두 생성의 매 단계에서 허용되는 토큰 집합을 제한해요. TensorRT-LLM에서는 사용자 정의 logits 후처리기를 정의해 현재 생성 단계에서 절대 쓰지 말아야 할 logits를 마스킹할 수 있습니다.

TensorRT-LLM 모델을 python 백엔드로 배포했을 때(즉 tensorrt_llm/config.pbtxt에서 triton_backend가 python으로 설정된 경우, Triton의 python 백엔드는 model.py로 TensorRT-LLM 모델을 서빙함), 커스텀 logits 처리기는 Executor 설정의 일부로 모델 초기화 시 logits_post_processor_map으로 지정해야 해요. 참고용 샘플입니다.

...

+ executor_config.logits_post_processor_map = {
+            "<custom_logits_processor_name>": custom_logits_processor
+           }
self.executor = trtllm.Executor(model_path=...,
                                model_type=...,
                                executor_config=executor_config)
...

또한 logits 후처리기를 요청별로 개별적으로 켜고 싶다면, 추가 input 파라미터로 할 수 있어요. 예를 들어 이 튜토리얼에서는 inflight_batcher_llm/tensorrt_llm/config.pbtxt에 logits_post_processor_name을 추가해요.

input [
  {
    name: "input_ids"
    data_type: TYPE_INT32
    dims: [ -1 ]
    allow_ragged_batch: true
  },
  ...
  {
    name: "lora_config"
	data_type: TYPE_INT32
	dims: [ -1, 3 ]
	optional: true
	allow_ragged_batch: true
- }
+ },
+ {
+   name: "logits_post_processor_name"
+   data_type: TYPE_STRING
+   dims: [ -1 ]
+   optional: true
+ }
]
...

그리고 inflight_batcher_llm/tensorrt_llm/1/model.py의 execute 함수에서 처리합니다.

def execute(self, requests):
    """`execute` must be implemented in every Python model. `execute`
    function receives a list of pb_utils.InferenceRequest as the only
    argument. This function is called when an inference is requested
    for this model.
    Parameters
    ----------
    requests : list
      A list of pb_utils.InferenceRequest
    Returns
    -------
    list
      A list of pb_utils.InferenceResponse. The length of this list must
      be the same as `requests`
    """
    ...

    for request in requests:
        response_sender = request.get_response_sender()
        if get_input_scalar_by_name(request, 'stop'):
            self.handle_stop_request(request.request_id(), response_sender)
        else:
            try:
                converted = convert_request(request,
                                            self.exclude_input_from_output,
                                            self.decoupled)
                ...

컨테이너를 시작할 때 tutorials 폴더를 /tutorials에 마운트하세요.

제공된 클라이언트 스크립트는 pydantic을 쓰므로 먼저 설치하세요.

pip install pydantic

옵션 1. 제공되는 클라이언트 스크립트 사용

먼저 JSON 답변 형식을 강제하지 않고 표준 요청을 보내 볼게요.

python3 /tutorials/AI_Agents_Guide/Constrained_Decoding/artifacts/client.py --prompt "Who is Harry Potter?" -o 100

다음과 같은 응답을 기대할 수 있어요.

Who is Harry Potter? Harry Potter is a fictional character in a series of fantasy novels written by British author J.K. Rowling. The novels chronicle the lives of a young wizard, Harry Potter, and his friends Hermione Granger and Ron Weasley, all of whom are students at Hogwarts School of Witchcraft and Wizardry. The main story arc concerns Harry's struggle against Lord Voldemort, a dark wizard who intends to become immortal, overthrow the wizard governing body known as the Ministry of Magic and subjugate all wizards and

이제 요청에 logits_post_processor_name을 지정해 볼게요.

python3 /tutorials/AI_Agents_Guide/Constrained_Decoding/artifacts/client.py --prompt "Who is Harry Potter?" -o 100 --logits-post-processor-name "outlines"

이번에는 다음과 같은 응답을 기대할 수 있어요.

Who is Harry Potter?{ "name": "Harry Potter","house": "Gryffindor","blood_status": "Pure-blood","occupation": "Wizards","alive": "No","wand": {"wood": "Holly","core": "Phoenix feather","length": 11 }}

보시다시피 utils.py에 정의한 스키마가 지켜졌어요. 참고로 LM Format Enforcer는 LLM이 생성 필드의 순서를 제어하게 하므로 필드 순서가 바뀌는 게 허용됩니다.

옵션 2. generate 엔드포인트 사용

먼저 JSON 답변 형식을 강제하지 않고 표준 요청을 보내 볼게요.

curl -X POST localhost:8000/v2/models/ensemble/generate -d '{"text_input": "Who is Harry Potter?", "max_tokens": 100, "bad_words": "", "stop_words": "", "pad_id": 2, "end_id": 2}'

다음과 같은 응답을 기대할 수 있어요.

{"context_logits":0.0,...,"text_output":"Who is Harry Potter? Harry Potter is a fictional character in a series of fantasy novels written by British author J.K. Rowling. The novels chronicle the lives of a young wizard, Harry Potter, and his friends Hermione Granger and Ron Weasley, all of whom are students at Hogwarts School of Witchcraft and Wizardry. The main story arc concerns Harry's struggle against Lord Voldemort, a dark wizard who intends to become immortal, overthrow the wizard governing body known as the Ministry of Magic and subjugate all wizards and"}

이제 요청에 logits_post_processor_name을 지정할게요.

curl -X POST localhost:8000/v2/models/ensemble/generate -d '{"text_input": "Who is Harry Potter?", "max_tokens": 100, "bad_words": "", "stop_words": "", "pad_id": 2, "end_id": 2, "logits_post_processor_name": "outlines"}'

이번에는 다음과 같은 응답을 기대할 수 있어요.

{"context_logits":0.0,...,"text_output":"Who is Harry Potter?{ \"name\": \"Harry Potter\",\"house\": \"Gryffindor\",\"blood_status\": \"Pure-blood\",\"occupation\": \"Wizards\",\"alive\": \"No\",\"wand\": {\"wood\": \"Holly\",\"core\": \"Phoenix feather\",\"length\": 11 }}"}

이렇게 하면 모델을 파인튜닝하지 않아도, 혹은 파인튜닝된 모델이어도 로직 기반으로 출력 형식을 정확히 강제할 수 있어요. 애플리케이션에서 스키마 준수가 필수라면 제약 디코딩을 워크플로에 꼭 활용해 보세요.