Ditto in Action: Torch-TensorRT로 손쉽게 LLM 추론 가속하기

Ditto를 소개합니다

Direct Torch to TensorRT-LLM Optimizer, 줄여서 Ditto는 Transformer 모델을 TensorRT 엔진으로 손쉽게 변환해주는 오픈소스 프레임워크입니다. 이 프로젝트는 SqueezeBits에서 개발해 공개했으며, Torch-TensorRT 라이브러리를 기반으로 PyTorch로 작성된 LLM을 TensorRT-LLM과 완벽하게 호환되는 TensorRT 엔진으로 바꿔 줍니다. 지난 6월에는 NVIDIA와 HPE Korea가 함께한 오픈 해커톤을 통해 새로운 기능을 추가하는 등, SqueezeBits에서 꾸준히 관리하고 있습니다.

Ditto는 Torch-TensorRT, TensorRT-LLM, 그리고 Transformers 라이브러리를 하나로 묶은 통합 프레임워크입니다. 덕분에 사용자는 PyTorch의 직관적인 인터페이스와 TensorRT의 강력한 추론 성능, 그리고 TensorRT-LLM의 LLM 특화 최적화 기능을 동시에 누릴 수 있습니다. 또한, TP, PP, LoRA, Quantization, MoE 등 최신 LLM 모델 서빙에 필수적인 기능도 Ditto 하나로 간편하게 적용할 수 있도록 지원합니다.

Ditto의 동작 원리

Ditto는 Torch-TensorRT를 기반으로 구현되어 PyTorch 연산을 TensorRT 연산으로 자동 변환할 수 있습니다. 이를 통해 Transformers 등 기존 라이브러리에 정의된 다양한 LLM을 코드 수정 없이 고성능 TensorRT 엔진으로 곧바로 변환할 수 있습니다.

PyTorch로 정의된 LLM 모델을 TensorRT 엔진으로 변환하는 과정은 다음 세 단계로 이뤄집니다.

1. 연산 그래프 추출: TorchDynamo 기반 컴파일러 스택을 활용하여 PyTorch 연산 그래프를 추출합니다.
2. 그래프 최적화: 추출된 연산 그래프를 TensorRT에 최적화된 형태로 재구성하며, 어텐션 연산 최적화 등의 LLM 특화 최적화 기법을 적용합니다.
3. 연산 변환 및 엔진 컴파일: 최적화된 그래프의 연산을 TensorRT 연산으로 변환한 뒤, TensorRT-LLM 호환 엔진으로 컴파일합니다.

Ditto는 PyTorch 컴파일러 스택 기반의 범용적인 접근을 통해 다양한 Transformer 기반 LLM 모델들을 빠르고 효율적으로 변환하며, 별도 코드 수정 없이 즉시 배포할 수 있도록 지원합니다. 모델에 구애받지 않는 이러한 접근 방식은 다양한 구조의 비전 인코더를 포함한 여러 Vision Language Model (VLM)도 손쉽게 변환하고, 실제 환경에 최적화하여 안정적으로 배포할 수 있도록 설계되었습니다.

PyTorch 모델에서 연산 그래프 추출하기

TensorRT 엔진 변환 과정은 연산 그래프 추출 단계로 시작합니다. Ditto는 PyTorch의 TorchDynamo 컴파일러 스택을 활용해 PyTorch로 정의된 LLM 모델에서 연산 그래프를 추출합니다. 이 과정에서 모든 텐서 연산은 PyTorch의 저수준 표현인 ATen 연산 단위로 분해되며, 이는 이후 최적화 및 변환 작업을 위한 토대가 됩니다.

연산 그래프 최적화

추출된 연산 그래프는 분석을 통해 다양한 그래프 최적화 기법이 적용된 형태로 변환됩니다. 여기에는 레이어 병합과 같은 범용 최적화뿐 아니라, 거대 언어 모델에 특화된 여러 최적화 기법이 포함되어, 이후 생성될 TensorRT 엔진이 최상의 성능을 발휘할 수 있도록 합니다.

이 단계에서 가장 집중적으로 최적화하는 부분은 Transformer 연산 중 핵심인 어텐션(attention) 계층입니다. Ditto는 LLM 모델 코드에 구현된 다양한 어텐션 방식을 정확히 식별하기 위해 패턴 인식 기반 그래프 최적화 기술을 활용합니다.

먼저 ATen 수준으로 분해된 어텐션 연산 패턴을 찾아내 임시 노드로 치환한 뒤, 다음 단계에서 이를 TensorRT 플러그인 노드로 대체합니다. 이를 통해 실행 시점에는 각 모델의 어텐션 연산이 TensorRT-LLM에서 제공하는 최적화된 어텐션 커널을 통해 수행되도록 변환합니다.

연산 그래프 변환

최적화된 연산 그래프는 마지막으로 Torch-TensorRT를 통해 TensorRT 엔진으로 변환됩니다. Torch-TensorRT는 PyTorch 모델을 NVIDIA GPU에서 최적화된 성능으로 실행할 수 있게 해 주는 NVIDIA의 공식 통합 라이브러리입니다. Torch-TensorRT에 관한 더욱 자세한 정보는 Accelerating Inference Up to 6x Faster in PyTorch with Torch-TensorRT 블로그에서 확인할 수 있습니다.

실행

Ditto를 통해 생성된 TensorRT 엔진은 TensorRT-LLM과 완벽하게 호환되므로, Python LLM API를 포함한 모든 기능을 별도의 수정 없이 그대로 활용할 수 있습니다. 따라서 Ditto로 변환한 엔진은 TensorRT-LLM에서 제공하는 추론 스크립트를 그대로 사용해 실행할 수 있습니다.

예시: Qwen2.5 모델군

이번 블로그에서는 Qwen2.5 모델군(Qwen2.5-7B-Instruct, Qwen2.5-VL-7B-Instruct, Qwen2.5-7B-Instruct-AWQ)을 다양한 프레임워크로 실행하고 성능을 비교하는 간단한 실험 결과를 소개합니다. 자세한 내용과 전체 성능 측정 스크립트는 Ditto의 GitHub 리포지토리를 참고하시기 바랍니다.

설치 및 준비

예시 벤치마크를 직접 실행하기 위해서는 다음과 같은 환경이 필요합니다.

• 27.1.1 이상 버전의 Docker가 설치된 환경
• nvcr.io/nvidia/tensorrt:25.05-py3 도커 이미지

NGC 카탈로그를 참고하여 nvcr.io/nvidia/tensorrt:25.05-py3 이미지로 도커 컨테이너를 생성한 뒤, 컨테이너 내 Bash 터미널에서 다음 커맨드로 Ditto를 설치하면 준비가 완료됩니다.

pip install git+https://github.com/SqueezeBits/Torch-TRTLLM.git@blog-post-blackwell

이 커맨드는 Ditto의 벤치마크 시점 브랜치와 더불어 아래의 패키지들을 포함한 모든 의존성 패키지들을 함께 설치합니다.

• PyTorch 2.7.1
• TensorRT 10.10.0.31
• Torch-TensorRT 2.6.1
• TensorRT-LLM 0.20.0

Torch-TensorRT와 TensorRT-LLM의 엄격한 버전 의존성으로 인해 상기한 패키지 일부와 설치 과정에서 충돌이 발생하나, Ditto가 각 패키지에서 활용하고 있는 부분에 대해서는 버전 의존성을 완화하여도 기능상 문제 없이 동작함을 확인하였습니다.

Naive PyTorch 모델 실행 예시

TensorRT-LLM에서는 기본적으로 어텐션 연산을 가속하기 위해 FlashAttention 방식이 반영된 최적화된 어텐션 커널을 사용합니다. 어텐션 커널의 성능 차이를 최소화하고 모델 변환으로 인한 성능 차이를 측정하기 위해 FlashAttention 라이브러리를 설치합니다.

pip install flash-attn==2.7.4.post1 --no-build-isolation

아래는 FlashAttention을 활용해 HuggingFace Transformers 모델의 성능을 평가하는 간소화된 Python 스크립트입니다.

import time
import torch
from transformers import AutoTokenizer, Qwen2ForCausalLM

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
DTYPE = torch.float16
BS = 32
ISL = 1024
OSL = 1024
MAGIC_STR = "pad"

if __name__ == "__main__":
    inputs = [MAGIC_STR * ISL] * BS
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    llm = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=DTYPE, attn_implementation="flash_attention_2").cuda()

    warmup_inputs = processor(texts[:min(2, BS)], return_tensors="pt")
    _ = llm.generate(**warmup_inputs.to("cuda"), max_new_tokens=output_len, min_new_tokens=output_len)

    torch.cuda.synchronize()
    st = time.perf_counter()
    inputs = processor(texts, return_tensors="pt")
    out = llm.generate(**inputs.to("cuda"), max_new_tokens=OSL, min_new_tokens=OSL)
    torch.cuda.synchronize()
    end = time.perf_counter()

    assert out.shape == (BS, ISL+OSL)

    latency_sec = end - st

    token_per_sec = (BS*OSL) / letency_sec
    print(f"LATENCY: {latency_sec} (sec)")
    print(f"THROUGHPUT: {token_per_sec} (token/sec)")

Ditto 실행 예시

아래는 Ditto를 사용해 변환된 TensorRT 엔진과 TensorRT-LLM LLMAPI를 활용해 LLM의 성능을 평가하는 Python 스크립트입니다.

import time

import torch
from transformers import AutoTokenizer, Qwen2ForCausalLM
from tensorrt_llm.llmapi import LLM, SamplingParams
from ditto.api import build_llm_engine, build_multimodal_engine

def main():
	MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
	DTYPE = torch.float16

	ENGINE_DIR = "./engines"

	BS = 32
	ISL = 1024
	OSL = 1024
	MAGIC_STR = "pad"

	build_llm_engine(
		Qwen2ForCausalLM.from_pretrained(model_id, torch_dtype=DTYPE),
		ENGINE_DIR,
		network_name=Qwen2ForCausalLM.__name__,
		max_batch_size=BS,
		max_seq_len=ISL+OSL,
	)
	
	inputs = [{"prompt": MAGIC_STR*ISL}]*BS
	tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
	sampling_params = SamplingParams(
		end_id=tokenizer.eos_token_id,
		max_tokens=OSL,
		ignore_eos=True
   )
	with LLM(ENGINE_DIR, tokenizer=tokenizer) as llm:
		warmup_inputs = inputs[:min(2, BS)]
		_ = llm.generate(warmup_inputs, sampling_params)

		torch.cuda.synchronize()
		st = time.perf_counter()
		output = llm.generate(inputs, sampling_params)
		torch.cuda.synchronize()
		end = time.perf_counter()

        assert all([len(o.prompt_token_ids) == input_len for o in output])
        assert all([len(o.outputs[0].token_ids) == output_len for o in output])
		
		latency_sec = end - st
		token_per_sec = (BS*OSL)/letency_sec
		print(f"LATENCY: {latency_sec} (sec)")
		print(f"THROUGHPUT: {token_per_sec} (token/sec)")
	
if __name__ == "__main__":
	main()

성능 측정 결과

실험은 NVIDIA B200 GPU 1개, 28개 vCPU, 283GB RAM을 갖춘 시스템에서 nvcr.io/nvidia/tensorrt:25.05-py3 이미지를 사용해 생성한 컨테이너 환경에서 진행되었습니다.

모든 지표에서 Ditto로 변환한 TensorRT 엔진은 기준이 되는 Hugging Face PyTorch 모델보다 일관되게 우수한 성능을 보여주었습니다. 이는 TensorRT 런타임의 효과적인 추론 가속 덕분이며, Ditto는 사용자가 다양한 LLM 모델에서 손쉽게 이러한 추론 가속을 손쉽게 구현할 수 있도록 지원합니다.

Ditto로 변환한 TensorRT 엔진에 대한 추가적인 성능 평가 결과, 배치 사이즈가 증가함에 따라 처리량이 예측 가능하게 확장되는 것으로 나타났습니다. 이러한 확장성 있는 성능과 유연성 덕분에 사용자는 여러 가지 배포 구성 중에서 선택하여 자신의 요구 사항에 맞게 성능을 조정 할 수 있습니다.

결론

Ditto는 직관적인 사용법과 간단한 설정 절차를 제공하여, 기존 LLM 배포 파이프라인에 쉽게 통합할 수 있습니다. 또한, NVIDIA GPU에서 최적화된 LLM 추론 성능을 달성하고, 다양한 모델에 걸쳐 최고 수준의 처리량과 지연 시간을 경험할 수 있도록 지원합니다. LLM을 가속할 준비가 되셨나요? 지금 바로 Ditto를 설치하고, 최고 성능의 LLM 서빙을 경험해 보세요.