LLM API 소개 (LLM API Introduction)
LLM API 소개 (LLM API Introduction)
TensorRT-LLM을 직접 다루려면 가장 먼저 만나게 되는 게 바로 LLM API예요. 이 고수준 Python API는 추론 워크플로를 한결 단순하게 만들어 주는데, 싱글 GPU부터 멀티 GPU·멀티 노드 배포까지 두루 지원하고 다양한 병렬화 전략과 고급 기능이 기본으로 깔려 있어요. 특히 겉은 단순해 보여도 속은 PyTorch 네이티브의 모듈형 백엔드라서, 커스터마이징이나 확장, 런타임 실험이 자유롭다는 게 매력적이에요.
빠른 시작 예제 (Quick Start Example)
TinyLlama로 첫 추론을 해 보는 가장 짧은 예제예요. 모델을 만들고, 프롬프트를 넣고, 생성 결과를 받아 보는 흐름이에요.
from tensorrt_llm import LLM, SamplingParams
def main():
# Model could accept HF model name, a path to local HF model,
# or Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF.
llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
# Sample prompts.
prompts = [
"Hello, my name is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
for output in llm.generate(prompts, sampling_params):
print(
f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
)
# Got output like
# Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
# Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the'
# Prompt: 'The capital of France is', Generated text: 'Paris.'
분산 추론, 멀티모달, 스펙큘레이티브 디코딩 같은 더 고급 사용법은 공식 저장소의 README에서 이어서 설명해 줘요.
모델 입력 (Model Input)
LLM() 생성자는 두 가지 입력을 받아요. Hugging Face 모델 ID를 쓰거나, 로컬 모델 경로를 넘기면 됩니다.
1. Hugging Face Hub에서 모델 가져오기
Hugging Face Model Hub에서 바로 가져오려면 저장소 이름(모델 ID)만 넘기면 돼요. 모델은 자동으로 내려받아집니다.
llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
NVIDIA가 제공하는 인기 모델의 양자화 체크포인트(FP4, FP8 등)도 같은 방식으로 사용할 수 있어요.
2. 로컬 Hugging Face 모델 사용하기
로컬 저장소의 모델을 쓰려면 먼저 직접 내려받아야 해요.
git lfs install
git clone https://huggingface.co/meta-llama/Meta-Llama-3.1-8B
그다음 로컬 디렉터리 경로를 지정해 모델을 불러옵니다.
llm = LLM(model="<local_path_to_model>")
참고: 일부 모델은 특정 라이선스 이용약관에 동의해야 해요. 내려받기 전에 약관에 동의하고 Hugging Face로 인증을 마쳐야 합니다.
팁과 트러블슈팅 (Tips and Troubleshooting)
TensorRT-LLM의 다른 API에 익숙한 새 사용자가 자주 겪는 문제를 정리해 두었어요.
RuntimeError: only rank 0 can start multi-node session, got 1
LLM API로 싱글 노드에서 멀티 GPU 추론을 할 때는 mpirun 접두사를 붙일 필요가 없어요. 예를 들어 python llm_inference_distributed.py처럼 그냥 실행하면 단일 노드에서 멀티 GPU를 사용합니다.
Slurm 노드에서 멈추는 문제
Slurm으로 관리하는 노드에서 멈추거나 문제가 생기면, 실행 스크립트 앞에 mpirun -n 1 --oversubscribe --allow-run-as-root 접두사를 붙여 보세요.
mpirun -n 1 --oversubscribe --allow-run-as-root python llm_inference_distributed.py
MPI_ABORT was invoked on rank 1 in communicator MPI_COMM_WORLD with errorcode 1.
LLM API가 mpi4py 라이브러리에 의존하기 때문에 생기는 문제예요. LLM 클래스를 함수 안에 두고, 프로그램의 진입점을 __main__ 네임스페이스 아래에 보호해 두면 재귀적 spawn을 피할 수 있어요.
docker run --net=host 사용 시 싱글 노드가 멈추는 문제
이 명령으로 관련 런타임 문제를 피할 수 있어요.
더 알아보기 (Learn more)
- 벤치마킹은
trtllm-serve로 실행하는 방법을 Run benchmarking with trtllm-serve에서 다뤄요. - 클래스별 상세한 시그니처는 API Reference에서 확인할 수 있어요.