Triton Inference Server

Triton Inference Server

NVIDIA Triton Inference Server의 임베딩 모델을 LiteLLM에서 사용하는 방법을 알아봐요.

출처: 문서

본문

LiteLLM은 Triton Inference Server의 임베딩 모델을 지원해요.

Triton /generate - 채팅 완성

triton 서버로 라우팅하려면 triton/ 접두사를 사용해요.

from litellm import completion
response = completion(
    model="triton/llama-3-8b-instruct",
    messages=[{"role": "user", "content": "who are u?"}],
    max_tokens=10,
    api_base="http://localhost:8000/generate",
)

config.yaml에 모델 추가:

model_list:
  - model_name: my-triton-model
    litellm_params:
      model: triton/"
      api_base: https://your-triton-api-base/triton/generate

Proxy 시작:

$ litellm --config /path/to/config.yaml --detailed_debug

OpenAI Python SDK로 요청:

import openai
from openai import OpenAI

# set base_url to your proxy server
# set api_key to send to proxy server
client = OpenAI(api_key="", base_url="http://0.0.0.0:4000")

response = client.chat.completions.create(
    model="my-triton-model",
    messages=[{"role": "user", "content": "who are u?"}],
    max_tokens=10,
)

print(response)

--header는 선택 사항이며, Virtual Keys와 함께 litellm proxy를 사용할 때만 필요해요.

curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer ***" \
--data ' {
"model": "my-triton-model",
"messages": [{"role": "user", "content": "who are u?"}]
}'

Triton /infer - 채팅 완성

triton 서버로 라우팅하려면 triton/ 접두사를 사용해요.

from litellm import completion

response = completion(
    model="triton/llama-3-8b-instruct",
    messages=[{"role": "user", "content": "who are u?"}],
    max_tokens=10,
    api_base="http://localhost:8000/infer",
)

config.yaml에 모델 추가:

model_list:
  - model_name: my-triton-model
    litellm_params:
      model: triton/"
      api_base: https://your-triton-api-base/triton/infer

Proxy 시작:

$ litellm --config /path/to/config.yaml --detailed_debug

OpenAI Python SDK로 요청:

import openai
from openai import OpenAI

# set base_url to your proxy server
# set api_key to send to proxy server
client = OpenAI(api_key="", base_url="http://0.0.0.0:4000")

response = client.chat.completions.create(
    model="my-triton-model",
    messages=[{"role": "user", "content": "who are u?"}],
    max_tokens=10,
)

print(response)

curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer ***" \
--data ' {
"model": "my-triton-model",
"messages": [{"role": "user", "content": "who are u?"}]
}'

Triton /embeddings - 임베딩

triton 서버로 라우팅하려면 triton/ 접두사를 사용해요.

from litellm import embedding
import os

response = await litellm.aembedding(
    model="triton/",
    api_base="https://your-triton-api-base/triton/embeddings", # /embeddings endpoint you want litellm to call on your server
    input=["good morning from litellm"],
)

config.yaml에 모델 추가:

model_list:
  - model_name: my-triton-model
    litellm_params:
      model: triton/"
      api_base: https://your-triton-api-base/triton/embeddings

Proxy 시작:

$ litellm --config /path/to/config.yaml --detailed_debug

OpenAI Python SDK로 요청:

import openai
from openai import OpenAI

# set base_url to your proxy server
# set api_key to send to proxy server
client = OpenAI(api_key="", base_url="http://0.0.0.0:4000")

response = client.embeddings.create(
    input=["hello from litellm"],
    model="my-triton-model"
)

print(response)

curl --location 'http://0.0.0.0:4000/embeddings' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer ***" \
--data ' {
"model": "my-triton-model",
"input": ["write a litellm poem"]
}'

더 알아보기 (Learn more)

• Triton Inference Server 공식 문서
• LiteLLM 임베딩 API