pytorch-backend

PyTorch (LibTorch) 백엔드

PyTorch용 Triton 백엔드는 PyTorch C++ API로 TorchScript 모델을 실행하도록 설계됐어요. Python API로 PyTorch에서 만든 모든 모델은 tracing 또는 scripting을 거쳐 TorchScript 모델로 만들어야 해요.

Triton 백엔드에 대한 더 자세한 내용은 Triton Backend 저장소에서, 질문·문제는 Triton Server 이슈에서 찾을 수 있어요. 아래 내용과 메인 서버 저장소의 일반 Triton 문서를 함께 읽어보는 걸 권해요.

참고: 이 문서의 대부분은 최신 AOT Inductor (PT2) 모델 아카이브 경로와 기존 TorchScript(LibTorch) 경로를 함께 다룬다는 점을 먼저 짚고 갈게요. 자기가 쓰는 형식에 해당하는 부분만 보면 돼요.

PyTorch 백엔드 빌드

최신 cmake로 빌드해요. 먼저 필수 의존성을 설치해요.

apt-get install rapidjson-dev python3-dev python3-pip
pip3 install patchelf==0.18.0

NVIDIA NGC 카탈로그의 적절한 PyTorch 컨테이너를 사용해야 해요. 예를 들어 NGC의 PyTorch 컨테이너 26.04 버전을 쓰는 백엔드를 빌드하려면:

mkdir build
cd build
cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DTRITON_PYTORCH_DOCKER_IMAGE="nvcr.io/nvidia/pytorch:26.04-py3" ..
make -j<N> install

빌드 과정에서 아래 필수 Triton 저장소가 내려와 쓰여요. 기본적으로 각 저장소의 main head가 쓰이지만, 나열된 CMake 인자로 값을 바꿀 수 있어요.

• triton-inference-server/backend: -DTRITON_BACKEND_REPO_TAG=[tag]
• triton-inference-server/core: -DTRITON_CORE_REPO_TAG=[tag]
• triton-inference-server/common: -DTRITON_COMMON_REPO_TAG=[tag]

커스텀 PyTorch로 빌드

현재 Triton은 PyTorch 백엔드에 특별히 패치된 PyTorch 버전을 요구해요. 이 버전들의 전체 소스는 NGC의 Docker 이미지로 제공돼요. 예를 들어 Triton 26.04 릴리스와 호환되는 PyTorch 버전은 nvcr.io/nvidia/pytorch:26.04-py3으로 제공되며 PyTorch 2.12.0a0를 지원해요.

추가 세부사항·버전 정보는 컨테이너 릴리스 노트에서 확인할 수 있어요.

PyTorch NGC 컨테이너에서 LibTorch와 TorchVision 헤더·라이브러리를 로컬 디렉터리로 복사해와요. 어떤 헤더/라이브러리를 복사해 오는지는 docker에서 확인할 수 있어요.

mkdir build
cd build
cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DTRITON_PYTORCH_INCLUDE_PATHS="<PATH_PREFIX>/torch;<PATH_PREFIX>/torch/torch/csrc/api/include;<PATH_PREFIX>/torchvision" -DTRITON_PYTORCH_LIB_PATHS="<LIB_PATH_PREFIX>" ..
make install

PyTorch 백엔드 사용하기

AOT Inductor 지원 (베타)

Triton 26.03 릴리스부터 ahead-of-time(AOT) Inductor로 컴파일·패키징된 모델 아카이브 지원이 생겼어요. 새 모델 아카이브 패키지(PT2)는 아래 예제처럼 만들 수 있으며, 보통 .pt2 확장자를 써요.

model.pt2 패키지 모델 아카이브를 만드는 예제 Python 코드:

ep = torch.export.export(model, sample_inputs)
torch._inductor.compile_and_package(ep, OUTPUT_FOLDER + "/model.pt2")

모델 저장소는 이렇게 보여야 해요:

model_repository/
`-- model_directory
    |-- 1
    |   `-- model.pt2
    `-- config.pbtxt

모델 구성

모델의 config.pbtxt에서 platform: "torch_aoti"를 주면 Triton이 PyTorch PT2 모델 아카이브에서 모델을 로드해요. Triton Server 26.05부터 Triton은 로드된 모델의 call specification(PyTorch 런타임이 제공)에 따라 런타임에 입출력 텐서 이름을 결정해요. 지원되는 명명 규칙은 두 가지예요.

순서 기반(Ordinal) I/O 이름 — 입출력을 call specification에서의 위치(0-based index)에 기반한 이름으로 다룰 수 있어요. 입력은 INPUT__ 뒤에 0-based 인덱스를 붙인 이름(INPUT__0), 출력은 OUTPUT__ 뒤에 0-based 인덱스를 붙인 이름(OUTPUT__0)을 사용해요.

Forward 기반 I/O 이름 — 모델의 forward 함수 명세와 컴파일할 때 쓴 샘플 입력에 기반한 이름으로 다룰 수 있어요. 입력은 torch.export.export가 처리하는 방식에 따라 둘로 나뉘어요.

• "arguments"는 항상 튜플로 취급되며 ARGS 프리픽스로 주소를 지정해요.
• "keyword arguments"는 항상 딕셔너리로 취급되며 KWARGS 프리픽스로 주소를 지정해요.

출력은 텐서, 튜플, 딕셔너리 중 하나일 수 있으며 RESULT 프리픽스로 주소를 지정해요. 리스트도 기술적으로 지원되지만 효과적으로 튜플처럼 취급돼요.

Forward 기반 I/O 이름의 기본 규칙

• 튜플 주소 지정은 ARGS[0], KWARGS[0], RESULT[0] 같은 Pythonic 튜플 의미론을 써요.
• 리스트 주소 지정은 튜플과 동일해요.
• 딕셔너리 주소 지정은 RESULT[key], KWARGS[key], ARGS[3][key] 같은 반-Pythonic 딕셔너리 의미론을 써요.
• 모델 출력이 단일 텐서면 접미사 접근자 없이 RESULT로 주소를 지정해요.

예제

import torch

SHAPE = (1, 16)

class ExampleAotiModel(torch.nn.Module):
  def __init__(self):
    super().__init__()

  def forward(
    self,
    hdata: torch.Tensor,
    vdata: torch.Tensor,
    options: dict[str, torch.Tensor],
    *
    kwag1: torch.Tensor,
    kw_arg2: torch.Tensor
  ) -> dict[str, torch.Tensor | tuple[torch.Tensor, torch.Tensor] | dict[str, torch.Tensor]]:
    return {
      # OUTPUT__0 or RESULT[AAA]
      "AAA": hdata + vdata,
      # OUTPUT__1 or RESULT[ZZZ]
      "ZZZ": hdata - vdata,
      "BBB": (
        # OUTPUT_2 or RESULT[BBB][0]
        hdata,
        # OUTPUT_3 or RESULT[BBB][1]
        vdata,
      ),
      # OUTPUT__4 or RESULT[CCC][keyA]
      # OUTPUT__5 or RESULT[CCC][key2]
      "CCC": options,
      "DDD": {
        # OUTPUT__6 or RESULT[DDD][KA]
        "KA": kwarg1,
        # OUTPUT__7 or RESULT[DDD][KB]
        "KB": kw_arg2,
      }
    }

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ExampleAotiModel()
model.to(device)
model = model.eval()

sample_args = (
  torch.zeros(SHAPE, dtype=torch.float16, device=device),  # INPUT__0 or ARGS[0]
  torch.zeros(SHAPE, dtype=torch.float16, device=device),  # INPUT__1 or ARGS[1]
  {                                                         # INPUT__2 or ARGS[2][keyA]
    "keyA": torch.zeros(SHAPE, dtype=torch.float16, device=device),
    "key2": torch.zeros(SHAPE, dtype=torch.float16, device=device),  # INPUT__3 or ARGS[2][key2]
  },
)
sample_kwargs = {
  "kwarg1": torch.zeros(SHAPE, dtype=torch.float16, device=device),  # INPUT__4 or KWARGS[kwarg1]
  "kw_arg2": torch.zeros(SHAPE, dtype=torch.float16, device=device), # INPUT__5 or KWARGS[kw_arg2]
}

exported_model = torch.export.export(model, args=sample_args, kwargs=sample_kwargs)
torch._inductor.aoti_compile_and_package(exported_model, package_path="model.pt2")

compiled_model = torch._inductor.aoti_load_package("model.pt2")
for spec in compiled_model.loader.get_call_spec():
  print(spec)

위 예제의 call specification을 통해 Triton은 다음을 알 수 있어요.

• 입력은 5개고 처음 2개는 이름이 없어서 ARGS[0], ARGS[1]으로 주소를 지정해요.
• 세 번째 입력은 keyA, key2 두 키를 가진 딕셔너리라 ARGS[2][keyA], ARGS[2][key2]로 주소를 지정해요.
• 네·다섯 번째 입력은 kwarg1, kw_arg2 키워드 인자라 KWARGS[kwarg1], KWARGS[kw_arg2]로 주소를 지정해요.
• 출력은 AAA, BBB, CCC, DDD, ZZZ 키를 가진 단일 딕셔너리예요.
  • AAA, ZZZ는 단순 텐서라 RESULT[AAA], RESULT[ZZZ]로 지정해요.
  • BBB는 원소 2개 튜플이라 RESULT[BBB][0], RESULT[BBB][1]로 지정해요.
  • CCC는 keyA, key2 키를 가진 딕셔너리라 RESULT[CCC][keyA], RESULT[CCC][key2]로 지정해요.
  • DDD는 KA, KB 키를 가진 딕셔너리라 RESULT[DDD][KA], RESULT[DDD][KB]로 지정해요.

이 지식을 바탕으로 Triton이 모델을 제대로 로드·상호작용하게 하는 config.pbtxt를 만들 수 있어요.

아래 예제는 모델 이름 방식이 어떻게 동작하는지 보여주기 위한 것으로, 간결하게 줄였어요. 이 구성은 불완전하므로 Triton Server에서 그대로 쓸 수 없어요.

backend: "pytorch"
platform: "torch_aoti"
input: [
  { name: "ARGS[0]" ... }, // "INPUT__0"
  { name: "ARGS[1]" ... }, // "INPUT__1"
  { name: "ARGS[2][keyA]" ... }, // "INPUT__2"
  { name: "ARGS[2][key2]" ... }, // "INPUT__3",
  { name: "KWARGS[kwarg1]" ... }, // INPUT__4,
  { name: "KWARGS[kw_arg2]" ... } // "INPUT__5"
]
output: [
  { name: "RESULT[AAA]" ... }, // "OUTPUT__0",
  { name: "RESULT[BBB][0]" ... }, // "OUTPUT__2",
  { name: "RESULT[BBB][1]" ... }, // "OUTPUT__3",
  { name: "RESULT[CCC][keyA]" ... }, // "OUTPUT__4",
  { name: "RESULT[CCC][key2]" ... }, // "OUTPUT__5",
  { name: "RESULT[DDD][KA]" ... }, // "OUTPUT__6",
  { name: "RESULT[DDD][KB]" ... }, // "OUTPUT__7",
  { name: "RESULT[ZZZ]" ... } // "OUTPUT__1"
]

중요: dict, tuple, list, torch.Tensor 타입만 지원되며, 모든 leaf(값) 타입은 반드시 torch.Tensor여야 해요. 다른 타입의 입출력은 지원되지 않아요. 딕셔너리 키에는 ,, [, ]를 포함할 수 없고 공백이나 출력 불가 문자도 포함할 수 없어요.

배칭: AOT Inductor 컴파일 모델은 기본(첫 번째 차원) 배칭이 지원돼요. 켜려면 다른 백엔드처럼 max_batch_size를 0보다 크게 설정하고 샘플별 dims(리딩 배치 차원 없이)를 지정해요. PT2 아카이브는 동적(dynamic) 첫(배치) 차원으로 export해야 해요. 그렇지 않으면 고정 배치 크기에 특화되어 런타임에 다른 배치 크기를 거부해요. torch.export.Dim을 torch.export.export의 dynamic_shapes 인자와 함께 써서 각 입력의 차원 0을 동적으로 표시하면 돼요.

시퀀스 배칭: 암시적 상태를 가진 시퀀스 배칭이 AOT Inductor 컴파일 모델에 지원돼요. 모델의 forward 함수는 시퀀스 제어 텐서(예: CONTROL_SEQUENCE_START, CONTROL_SEQUENCE_READY, CONTROL_SEQUENCE_CORRID)와 암시적 상태 입력을 일반 텐서 인자로 받고, 새 상태 값을 출력 중 하나로 반환해야 해요. CONTROL_SEQUENCE_CORRID는 숫자 data_type(예: TYPE_INT32, TYPE_INT64)으로 설정했을 때 지원돼요. TYPE_STRING correlation id는 AOT Inductor 런타임이 텐서 전용이라 지원되지 않아요.

경고 - 미지원: PT2 모델 아카이브로 패키징된 AOT Inductor 컴파일 모델에 다음 기능은 아직 미지원이에요. (1) TYPE_STRING 입출력(문자열 시퀀스 상태와 TYPE_STRING CONTROL_SEQUENCE_CORRID 포함 — AOT Inductor 런타임은 텐서 전용), (2) 복잡한(중첩 dict/tuple/list) 입출력을 가진 모델의 배칭 — 기본 배칭은 입출력이 평범한 텐서인 모델을 대상으로 해요, (3) RHEL(manylinux) 빌드 — AOT Inductor는 RHEL에서 미지원.

모델 초기화 훅 (MODEL_INIT_LIBRARY)

일부 모델은 model.pt2 안에서 표현할 수 없는, 로드 시(모델 패키지 로드 전) 1회성 네이티브 초기화를 실행해야 해요. 전형적인 경우는 커스텀 연산자가 실행 시점에 프로세스 전역 레지스트리에서 큰 가중치를 해석하는데, 그 가중치가 패키지 밖(예: model.pt2 옆 sidecar 파일)에 있어서 먼저 메모리로 읽어 등록해야 하는 모델이에요.

PT2 경로는 이를 위한 선택적·모델별 훅을 지원해요. config.pbtxt가 MODEL_INIT_LIBRARY 파라미터를 설정하면, 백엔드는 모델 로드 시 그 공유 라이브러리를 dlopen()하고 초기화 진입점을 호출해요. 백엔드는 그 라이브러리를 링크하지 않고 무엇을 하는지도 알지 못해요. 그냥 로드하고, 고정 진입점을 호출하고, 언로드 시 해제할 뿐이에요.

parameters {
  key: "MODEL_INIT_LIBRARY"
  value: { string_value: "/abs/path/to/libmy_model_init.so" }
}

라이브러리는 다음 C 진입점을 export해야 하며, 선택적으로 finalizer를 export할 수 있어요.

// Called once at model load, before the model package is loaded.
//   model_dir    : the versioned model directory (the one containing model.pt2)
//   device_index : the GPU ordinal for the instance, or -1 for CPU
// Returns an opaque handle that is passed back to the finalizer on unload, or
// NULL if there is nothing to keep.
extern "C" void* triton_pytorch_model_init(const char* model_dir, int device_index);

// Optional. Called once on model unload with the handle returned above.
extern "C" void  triton_pytorch_model_fini(void* state);

MODEL_INIT_LIBRARY가 설정되지 않으면(기본) 라이브러리가 로드되지 않고 완전히 no-op이에요. 설정됐지만 dlopen() 실패하거나 triton_pytorch_model_init을 export하지 않으면 모델 로드가 오류로 실패해요.

경고: MODEL_INIT_LIBRARY는 서버 프로세스에서, 서버 권한으로 모델 구성에 주어진 경로의 네이티브 코드를 로드·실행하게 해요. 신뢰하는 라이브러리만 가리켜야 해요. 이건 이미 모델 저장소 자체가 요구하는 신뢰 수준과 같아요 — PT2 패키지에는 로드될 때 실행되는 컴파일된 코드가 포함되므로 — 신뢰 경계를 넓히진 않지만, 모델 저장소는 신뢰되고 연산자에 의해 제어되는 소스로 남아 있어야 해요.

HSTU (Generative Recommenders)

HSTU(Hierarchical Sequential Transduction Unit)는 Generative Recommenders(GR)에서 쓰는 시퀀스 인코더예요. GR은 추천을 후보 항목을 독립적으로 점수화하는 문제가 아니라, 사용자 상호작용 이력 위의 순차적·생성적 모델링 문제로 재구성해요. HSTU 생성 추천 모델은 platform: "torch_aoti"로 서빙할 수 있어요. 종단간 가이드는 HSTU 튜토리얼과 recsys-examples HSTU 추론 문서를 참고해요.

PyTorch 2.0 모델

PyTorch 2.0 기능을 사용할 수 있어요. 다만 Triton의 PyTorch 백엔드는 model.pt 파일 형태의 직렬화된 모델 표현을 요구해요. 모델의 직렬화 표현은 PyTorch의 torch.save()로 만들어 model.pt 파일을 생성할 수 있어요. 새 PT2 모델 아카이브 포맷 지원은 개발 중이며 Triton 26.03 이상에서 테스트할 수 있어요.

model.pt 패키지 모델 아카이브를 만드는 예제 Python 코드:

traced = torch.jit.script(model)
traced.save(OUTPUT_FOLDER + "/model.pt")

모델 저장소는 이렇게 보여야 해요:

model_repository/
`-- model_directory
    |-- 1
    |   `-- model.pt
    `-- config.pbtxt

TorchScript 모델

모델 저장소는 위와 동일하게 생겼어요. model.pt가 TorchScript 모델 파일이에요.

구성 (Configuration)

Triton은 config.pbtxt의 Parameters 섹션을 통해 TorchScript 모델의 실행 모드를 제어하는 몇 가지 플래그를 노출해요.

구성 옵션

• default_model_name: Triton PyTorch 백엔드가 주어진 이름의 파일에서 모델을 로드하도록 지시해요.

default_model_name: "another_file_name.pt2"

• platform: AOT Inductor 또는 기존 LibTorch 프레임워크 중 어느 것으로 모델을 로드할지 지시해요. 기존 LibTorch는 pytorch_libtorch, AOT Inductor는 torch_aoti.

platform: "torch_aoti"

파라미터

• DISABLE_OPTIMIZED_EXECUTION: TorchScript 모델의 최적화된 실행을 끄는 불리언 플래그예요. 기본적으로 최적화 실행은 항상 켜져 있어요. 로드된 TorchScript 모델에 대한 초기 호출은 상당한 시간이 걸려요. 이런 긴 모델 워밍업 때문에 Triton은 최적화 없이도 실행을 허용해요. 일부 모델에서는 최적화 실행이 성능에 도움이 안 되거나 오히려 나쁜 영향을 주기도 해요.

parameters: {
  key: "DISABLE_OPTIMIZED_EXECUTION"
  value: { string_value: "true" }
}

• INFERENCE_MODE: TorchScript 모델의 Inference Mode 실행을 켜는 불리언 플래그예요. 기본적으로 inference mode는 켜져 있어요. InferenceMode는 NoGradMode와 유사한 새 RAII 가드로, 연산이 autograd와 전혀 상호작용하지 않는다고 확신할 때 쓰여요. NoGradMode와 달리 이 모드에서 실행되는 코드는 autograd를 비활성화해 더 나은 성능을 얻어요. 일부 모델에서는 성능에 도움이 안 되고 드물게 부정적일 수 있다는 점을 기억해 두세요.

parameters: {
  key: "INFERENCE_MODE"
  value: { string_value: "true" }
}

• DISABLE_CUDNN: cuDNN 라이브러리를 끄는 불리언 플래그예요. 기본적으로 cuDNN은 켜져 있어요. cuDNN은 심층 신경망의 GPU 가속 프리미티브 라이브러리로, 표준 루틴에 고도로 튜닝된 구현을 제공해요. 보통 cuDNN을 켜면 모델이 더 빠르게 실행되지만, cuDNN이 더 느리거나 메모리 사용이 높거나 오류를 일으키는 예외도 있어요.

parameters: {
  key: "DISABLE_CUDNN"
  value: { string_value: "true" }
}

• ENABLE_WEIGHT_SHARING: 같은 장치의 모델 인스턴스가 가중치를 공유하도록 켜는 불리언 플래그예요. 이 최적화는 stateful 모델에 쓰면 안 돼요. 지정하지 않으면 가중치 공유는 꺼져 있어요.

parameters: {
  key: "ENABLE_WEIGHT_SHARING"
  value: { string_value: "true" }
}

• ENABLE_CACHE_CLEANING: 각 모델 실행 후 CUDA 캐시 정리를 켜는 불리언 플래그예요. 지정하지 않으면 꺼져 있어요. 모델이 CPU에 있으면 효과가 없어요. true로 설정하면 매 실행마다 추가 CUDA 캐시 정리 연산 때문에 성능에 부정적일 가능성이 높아요. 그래서 여러 모델을 서빙하면서 실행 중 CUDA out-of-memory 문제를 겪을 때만 써야 해요.

parameters: {
  key: "ENABLE_CACHE_CLEANING"
  value: { string_value: "true" }
}

• INTER_OP_THREAD_COUNT: PyTorch는 TorchScript 모델 추론에 여러 CPU 스레드를 쓸 수 있어요. 하나 이상의 추론 스레드가 주어진 입력에 대해 모델의 forward pass를 실행해요. 각 추론 스레드는 모델의 연산을 하나씩 인라인으로 실행하는 JIT 인터프리터를 호출해요. 이 파라미터가 그 스레드풀의 크기를 설정해요. 기본값은 CPU 코어 수예요.

parameters: {
  key: "INTER_OP_THREAD_COUNT"
  value: { string_value: "1" }
}

이 파라미터는 PyTorch 백엔드에 전역으로 설정돼요. 이 파라미터를 지정한 첫 번째 모델 구성 파일의 값이 쓰이며, 이후 다른 모델 구성 파일의 값(다르다면)은 무시돼요.

• INTRA_OP_THREAD_COUNT: inter-op 병렬성에 더해, PyTorch는 연산 내부에서도 여러 스레드를 쓸 수 있어요(intra-op parallelism). 큰 텐서의 요소별 연산, 합성곱, GEMM, 임베딩 조회 등에서 유용할 수 있어요. 기본값은 CPU 코어 수예요.

parameters: {
  key: "INTRA_OP_THREAD_COUNT"
  value: { string_value: "1" }
}

• 추가 최적화: 특정 Torch 최적화를 끌 수 있는 불리언 파라미터 세 개가 더 있어요 — ENABLE_JIT_EXECUTOR, ENABLE_JIT_PROFILING, ENABLE_JIT_EXECUTOR 관련. 이들은 복잡한 실행 모드와 동적 shape를 가진 모델에서 지연 회귀를 일으킬 수 있는 최적화를 끄는 데 쓰여요. 지정하지 않으면 기본적으로 모두 켜져 있어요.

모델 인스턴스 그룹 종류 (Model Instance Group Kind)

PyTorch 백엔드는 입력 텐서 배치 방식에 따라 다음 인스턴스 그룹 종류를 지원해요.

• KIND_GPU: 입력이 모델 인스턴스와 연결된 GPU 장치에서 준비돼요.
• KIND_CPU: 입력이 CPU에서 준비돼요.
• KIND_MODEL: 입력이 CPU에서 준비돼요. 모델 로드 시 백엔드가 GPU 장치를 고르지 않고, 모델에 지정된 장치를 존중해 그대로 추론에 써요. 모델이 내부적으로 여러 GPU를 쓸 때 유용해요.

중요: 모델에 장치가 지정되지 않으면 백엔드는 첫 번째 사용 가능한 GPU 장치를 사용해요.

instance_group {
   count: 2
   kind: KIND_GPU
}

커스터마이징 (Customization)

다음 PyTorch 설정은 config.pbtxt에 파라미터를 설정해 커스터마이즈할 수 있어요.

중요: 아래 옵션은 platform: "pytorch_libtorch"를 쓸 때만 적용돼요.

• NUM_THREADS: torch.set_num_threads(int) — CPU에서 intra-op 병렬화에 쓰는 스레드 수.

parameters: {
  key: "NUM_THREADS"
  value: { string_value: "4" }
}

• NUM_INTEROP_THREADS: torch.set_num_interop_threads(int) — CPU에서 interop 병렬화(JIT 인터프리터 등)에 쓰는 스레드 수.
• TORCH_COMPILE_OPTIONAL_PARAMETERS: torch.compile() 파라미터를 JSON 객체로 인코딩한 것. fullgraph(bool), dynamic(bool), backend(str), mode(str: "default", "reduce-overhead", "max-autotune"), options(dict), disable(bool).

parameters: {
  key: "TORCH_COMPILE_OPTIONAL_PARAMETERS"
  value: { string_value: "{\"disable\": true}" }
}

경고: torch.compile로 최적화 가능한 Python 함수는 model.py 파일에서 직접 서빙할 수 없어요. torch.nn.Module을 상속하는 클래스로 감싸야 해요. 모델 가중치는 같은 GPU 장치의 여러 인스턴스가 공유할 수 없어요.

중요 참고 사항 (Important Notes)

• GPU에서 PyTorch 모델 실행은 본질적으로 비동기예요. 그래서 PyTorch 모델 실행 오류가 이후 몇 번의 추론 요청에서 발생할 수 있어요. 서버 시작 시 환경 변수 CUDA_LAUNCH_BLOCKING=1을 설정하면 동기 실행을 강제해 실패 사례를 제대로 디버깅할 수 있어요. 그런 경우 모델이 실패 상태에서 회복될 수도 있고 아닐 수도 있어서, 서빙을 계속하려면 서버 재시작이 필요할 수 있어요.
• PyTorch는 String Tensor를 지원하지 않지만, List of String을 입출력으로 받는 모델은 지원해요. 이런 모델에 대해 Triton은 String 데이터타입을 써서 String 입출력을 전달하게 해요. String I/O에 List를 Tensor 대신 쓰는 제약 때문에, String 타입 I/O는 1-dim 입출력만 지원돼요.
• 멀티 GPU 환경에서 Tracing으로 TorchScript 모델을 만들면 잠재적 런타임 문제가 생길 수 있어요. 이 문제는 모델 인스턴스와 텐서 사이의 장치 불일치 때문에 생겨요. 기본적으로 Triton은 사용 가능한 각 GPU마다 모델 실행 인스턴스를 하나씩 만들어요. TorchScript 생성 과정에서 쓴 GPU 장치와 다른 GPU 장치의 모델 인스턴스로 요청이 보내지면 런타임 오류가 발생해요. 멀티 GPU에서는 Tracing 대신 Scripting을 쓰는 걸 강력히 권해요. Tracing이 불가피하면 모델 구성에서 모델 인스턴스의 GPU 장치를 명시적으로 지정해, 모델을 tracing한 GPU와 같은 장치에 모델 인스턴스와 추론 텐서가 배정되게 하면 돼요.
• KIND_MODEL을 모델 인스턴스 종류로 쓰면, 모델의 첫 번째 파라미터의 기본 장치가 사용돼요.

---

출처: 공식문서 (PyTorch Backend)