Related Resources
Agentic AI / Generative AI

수학 문제 해결 속도를 4배 높이는 추론 가속 전략

2025년 11월 14일
By , , and
Discuss (0)
Reading Time: 5 minutes

거대 언어 모델은 뛰어난 문제 해결 능력을 갖추고 있지만, 실제 서비스 환경에서 빠르고 안정적으로 활용하려면 성능만으로는 부족합니다. 제대로 작동하려면 서빙 스택, 양자화 방식, 디코딩 전략까지 모두 갖춰야 하고, 이 구성 요소들은 대부분 따로따로 존재해 통합이 까다롭습니다. 이 때문에 팀은 다양한 설정을 실험하기 위해 여러 컨테이너, 변환 스크립트, 임시 코드 등을 동시에 다뤄야 하는 번거로움을 겪게 됩니다.

이 글에서는 NVIDIA NeMo-Skills 라이브러리를 활용해 NVIDIA TensorRT-LLM을 효율적으로 운영하고, 빠르고 재현 가능한 추론 파이프라인을 구축하는 방법을 소개합니다. 여기서 소개하는 구성은 AI Mathematical Olympiad Prize 2024 수상 모델을 기반으로 하며, FP8 양자화와 ReDrafter 기반 speculative decoding을 적용해 NVIDIA H100 GPU 2개 환경에서 배치 추론 속도를 4배까지 끌어올린 사례를 단순화한 버전입니다. 동일한 워크플로우는 단일 워크스테이션에서도 실행 가능하고, 소규모 수정만으로 클러스터 환경까지 확장할 수 있습니다.

이 글을 따라가며 아래 내용을 직접 실습해볼 수 있습니다:

  • OpenMath 모델을 FP8 TensorRT-LLM 엔진으로 변환 및 준비하는 방법
  • speculative decoding을 위한 ReDrafter draft 모델을 학습하고 통합하는 방법
  • 보안 코드 샌드박스를 포함한 최적화된 추론 서버 실행 방법
  • BF16, FP8, FP8+ReDrafter 구성별 지연 시간 및 처리량 벤치마크 방법

실습 시, H100 GPU 2개 또는 FP8을 지원하는 동급 GPU가 장착된 머신, 혹은 유사한 사양의 Slurm 클러스터 환경이 권장됩니다.

환경 설정하기

먼저 일관된 실행 환경을 만들기 위해 NVIDIA PyTorch NGC 컨테이너를 사용합니다. 여기에 모델 최적화를 위한 TensorRT-LLM과 파이프라인 관리를 위한 NeMo-Skills 라이브러리를 설치합니다. FP8 추론을 사용하려면 FP8을 지원하는 GPU가 필요하며, NVIDIA Ada Lovelace, Hopper, Blackwell, Rubin 아키텍처가 이에 해당합니다. 이 글에서는 GPU 2개가 있는 환경을 기준으로 설명합니다.

컨테이너 실행 및 라이브러리 설치

nvcr.io/nvidia/pytorch:25.05-py3 컨테이너에 진입한 뒤, 아래 명령어로 TensorRT-LLM과 NeMo-Skills를 설치합니다:

# Ensure no conflicting TensorRT installations and install TensorRT-LLM
[ -f /etc/pip/constraint.txt ] && : > /etc/pip/constraint.txt
pip uninstall -y tensorrt
pip3 install tensorrt_llm==1.1.0rc0
 
# Install NeMo-Skills
pip install git+https://github.com/NVIDIA/NeMo-Skills.git

모델 가중치 준비하기

이제 거대 언어 모델(LLM)을 준비할 차례입니다. nvidia/OpenMath-Nemotron-14B-Kaggle 모델을 다운로드한 후, FP8 양자화를 적용해 최적화된 TensorRT-LLM 엔진으로 변환합니다.

※ FP8 양자화 참고: FP8(Floating Point 8-bit)은 매우 효율적인 양자화 방식이지만, E4M3 FP8을 지원하는 GPU(Hopper 아키텍처 등)에서만 사용할 수 있습니다. 그 외 GPU 환경에서는 calibration이 필요 없는 int8_wo(가중치 전용 8비트 정수 양자화)를 사용하는 것이 적합합니다.

모델 가중치 및 데이터셋 다운로드

먼저 Hugging Face 토큰을 발급받아 환경 변수로 등록한 뒤, Hugging Face CLI를 사용해 필요한 모델과 데이터셋을 다운로드합니다.

# Export your Hugging Face token
export HF_TOKEN=hf_YOUR_HUGGING_FACE_TOKEN 
 
# Install Hugging Face CLI
pip install -U "huggingface_hub[cli]"
 
# Download the 14B parameter main model
huggingface-cli download nvidia/OpenMath-Nemotron-14B-kaggle --local-dir OpenMath-Nemotron-14B-kaggle
 
# Download the OpenMathReasoning dataset for calibration
huggingface-cli download nvidia/OpenMathReasoning --repo-type dataset --local-dir OpenMathReasoning

FP8 양자화를 위한 캘리브레이션 데이터셋 준비

FP8 양자화를 적용하려면, 실제 추론에 사용될 입력 데이터를 대표할 수 있는 소규모 캘리브레이션 데이터셋이 필요합니다. 여기서는 OpenMathReasoning 데이터셋의 일부를 활용해 해당 데이터를 생성합니다. Hugging Face 형식으로 캘리브레이션용 수학 데이터셋을 생성하는 예제가 함께 제공됩니다.

TensorRT-LLM 엔진 변환 및 양자화

이제 Hugging Face 모델을 TensorRT-LLM 엔진으로 변환하고, 앞서 준비한 캘리브레이션 데이터셋을 활용해 FP8 양자화를 적용합니다. 이 과정을 통해 FP8 추론에 최적화된 LLM 엔진이 생성됩니다.

ns convert \
    --input_model OpenMath-Nemotron-14B-kaggle \
    --output_model OpenMath-Nemotron-14B-kaggle-fp8-trtllm \
    --convert_from hf \
    --convert_to trtllm \
    --num_gpus 2 \
    --dtype fp8 \
    --hf_model_name nvidia/OpenMath-Nemotron-14B-kaggle \
    --model_type qwen \
    --max_input_len 30000 \
    --max_seq_len 32000 \
    --no-trt_reuse_tmp_engine \
    --calib_dataset ./calibration_dataset

이 명령어를 실행하면 FP8로 양자화된 LLM 엔진이 배포 준비를 마치게 됩니다.

ReDrafter로 추론 속도 더 끌어올리기

추론 효율을 한층 더 끌어올리기 위해 ReDrafter를 통합합니다. ReDrafter는 작은 크기의 “draft” 모델이 먼저 토큰을 예측하고, 메인 LLM이 이를 확인하는 방식의 speculative decoding 기법입니다. 이 방식은 전체 응답 생성을 더 빠르게 만들어줍니다. ReDrafter는 Apple이 개발한 RNN 기반 추론 방식이며, TensorRT-LLM에서 지원하는 대부분의 모델과 호환됩니다.

ReDrafter 설치 및 학습

먼저 ReDrafter 라이브러리를 설치합니다. draft 모델의 토크나이저와 학습 데이터는 베이스 모델과 동일해야 하며, 원본 학습 데이터가 없다면 베이스 모델의 출력 결과를 활용해 draft 모델을 학습시킬 수도 있습니다.

# Install the ReDrafter library
pip install --no-binary=protobuf --ignore-requires-python \     
"git+https://github.com/apple/ml-recurrent-drafter.git#egg=recurrent-drafting[dev,train]"
 
# Train the ReDrafter model
ns run_cmd --log_dir ./logs/ \
torchrun --nproc_per_node=2 -m nemo_skills.training.train_redrafter \
    --llm_name_or_path 'OpenMath-Nemotron-14B-kaggle' \
    --dataset "OpenMathReasoning" \
    --dataset_split "tir" \
    --bf16 True \
    --output_dir "redrafter_output" \
    --num_train_epochs 1 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --save_strategy "no" \
    --learning_rate 0.001 \
    --weight_decay 0. \
    --warmup_ratio 0.1 \
    --lr_scheduler_type "cosine" \
    --logging_steps 20 \
    --tf32 True \
    --model_max_length 2048 \
    --dataset_nrows 50000 \
    --drafter_predict_n_tokens 3 \
    --drafter_num_layers 2 \
    --rnn True \
    --phase train \
    --report_to wandb # Remove if not using wandb

학습 과정에서는 redrafter2_top1 스코어를 확인해야 합니다. 이 값이 0.6 이상이면, 전체 추론 단계 중 약 60%에서 초안 토큰 3개가 그대로 채택된다는 의미이며, 이는 실행 시간 기준으로 약 2배의 속도 향상을 기대할 수 있습니다.

ReDrafter 모델용 TensorRT-LLM 엔진 생성

이제 학습이 완료된 ReDrafter 모델을 TensorRT-LLM 체크포인트로 변환하고, 이를 메인 LLM과 결합해 최종 가속형 TensorRT-LLM 엔진을 생성합니다.

먼저, 변환 스크립트를 사용하기 위해 TensorRT-LLM 리포지토를 복제합니다:

git clone https://github.com/NVIDIA/TensorRT-LLM/

다음으로, 학습된 ReDrafter PyTorch 체크포인트를 TensorRT-LLM 포맷으로 변환합니다. 이 단계에서는 TensorRT-LLM에서 제공하는 변환 스크립트를 활용해 draft 모델을 추론 엔진에 통합 가능한 형태로 준비하게 됩니다.

# Base model intermediate checkpoint from FP8 quantization step
export BASE_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-trtllm-tmp-ckpt
# Trained draft checkpoint
export REDRAFTER_PYTORCH_CKPT=$(pwd)/redrafter_output/redrafter__redrafter_OpenMath-Nemotron-14B-kaggle_n_3_lr_0.001_layers_2
export REDRAFTER_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-draft-ckpt
 
cd ./TensorRT-LLM/examples/redrafter
python convert_checkpoint.py \
    --base_model_checkpoint_dir $BASE_TRTLLM_CKPT \
    --drafter_model_dir $REDRAFTER_PYTORCH_CKPT \
    --output_dir $REDRAFTER_TRTLLM_CKPT \
    --dtype bfloat16 \
    --tp_size 2 \
    --redrafter_num_beams 1 \
    --redrafter_draft_len_per_beam 3
cd ../../../

마지막으로, 메인 LLM과 draft 모델을 결합해 speculative decoding이 가능한 최종 TensorRT-LLM 엔진을 생성합니다. 이 엔진은 base 모델과 draft head를 하나로 통합하여, 추론 속도와 효율성을 극대화한 구성으로 완성됩니다.

trtllm-build \
    --checkpoint_dir $REDRAFTER_TRTLLM_CKPT \
    --output_dir OpenMath-Nemotron-14B-kaggle-fp8-redrafter-trtllm \
    --gemm_plugin fp8 \
    --use_paged_context_fmha=enable \
    --max_batch_size 32 \
    --max_seq_len 32000 \
    --max_input_len 32000 \
    --max_num_tokens 32000 \
    --speculative_decoding_mode explicit_draft_tokens \
    --max_beam_width 1 \
    --kv_cache_type paged

이제 ReDrafter가 결합된 TensorRT-LLM 엔진이 완성되었으며, 바로 추론 서버에 적용할 수 있습니다!

벤치마킹 및 결과

전체 파이프라인을 직접 실행해볼 수 있도록 별도 노트북도 함께 제공합니다. 이 노트북은 앞서 설명한 컨테이너 환경 및 라이브러리 설치와 동일한 구성에서 실행되었으며, H100 GPU 2개를 기반으로 추론을 수행합니다. 노트북에서는 다음과 같은 작업을 실습할 수 있습니다:

  • BF16, FP8, FP8+ReDrafter 구성별 TensorRT-LLM 엔진으로 추론 실행
  • 첫 토큰 생성 시간, 장비당 처리량 등 다양한 성능 지표 비교
  • 고급 제어 기능 사용: 고정 시간 이후 강제 종료, N개의 응답이 완료되면 종료 등
  • 툴 호출 기능을 포함한 추론 실행

아래는 노트북에서 확인할 수 있는 벤치마크 결과의 예시입니다:

지표BF16FP8FP8+ReDrafter
전체 생성 시간 (초)144.264.730.5
샘플당 평균 처리량 (토큰/초)34.675.2138.5

표 1. H100 GPU 2개 환경에서 구성별 TensorRT-LLM 성능 비교

전체 벤치마크 결과와 코드는 노트북에서 확인할 수 있으며, 더 많은 결과는 AIMO-2 수상 솔루션 논문에서도 살펴볼 수 있습니다.

선택: 툴 호출 및 코드 실행 샌드박스 활성화

OpenMath LLM은 단순한 텍스트 생성 모델을 넘어, 툴 사용 명령을 이해하고 실행할 수 있는 강력한 추론 모델입니다. 즉, 텍스트뿐만 아니라 Python 코드를 생성하고, 이를 보안된 샌드박스 환경에서 실행해 문제를 해결할 수 있습니다. 제공된 노트북에서는 LLM 서버와 코드 실행 샌드박스를 함께 실행하는 예제를 포함하고 있습니다.

작동 방식은 다음과 같습니다:

  1. LLM이 <tool_call>과 </tool_call> 토큰으로 감싼 Python 코드를 생성합니다.
  2. 추론 엔진이 해당 코드를 추출해 샌드박스로 전송합니다.
  3. 샌드박스가 코드를 실행하고 결과를 반환합니다.
  4. 이 결과는 LLM에 다시 전달되어 후속 응답 생성 또는 최종 답안 작성에 활용됩니다.

다음은 이러한 상호작용의 예시입니다:

<tool_call>
# Initialize a list to store valid bases
valid_bases = []
 
 
# Check bases from 10 upwards
for b in range(10, 10000):  # Arbitrary large upper limit
    num1 = 9 * b + 7
    num2 = b + 7
    if num1 % num2 == 0:
        valid_bases.append(b)
        print(f"Found base: {b}")
 
 
# Sum the valid bases
sum_bases = sum(valid_bases)
print(f"Sum: {sum_bases}")
 
 
# If sum is over 1000, take modulo 1000
if sum_bases > 1000:
    result = sum_bases % 1000
else:
    result = sum_bases
 
 
print(f"Final Result: {result}")
</tool_call>
```output
Found base: 21
Found base: 49
Sum: 70
Final Result: 70
```

툴 호출 기능을 끄고 싶다면, NeMo-Skills 문서에 나와 있는 대로 get_code_execution_model 대신 get_model을 사용하면 됩니다.

직접 실행해보세요. 노트북을 통해 여러분의 하드웨어에서 성능 향상을 벤치마크하고, 툴 호출 기능도 실험해볼 수 있습니다.

Discuss (0)

Tags

Agentic AI / Generative AI | Data Science | Developer Tools & Techniques | AI Inference | featured | Kaggle | LLMs

작성자 소개

Igor Gitman 프로필
Igor Gitman은 NVIDIA의 Applied Scientist로, 그의 연구는 주로 LLM의 수학적 추론 능력을 향상시키는 데 초점을 맞추고 있습니다. 또한 NVIDIA NeMo-Skills 라이브러리의 주요 유지 관리자 중 한 명이기도 합니다. Igor는 Carnegie Mellon University에서 머신러닝 전공으로 석사 학위를 받았습니다.
Rakib Hasan 프로필
Rakib Hasan은 NVIDIA의 Senior AI Developer Technology Engineer로, 대규모 언어 모델(LLM) 추론을 포함한 딥러닝 워크로드 최적화에 전문성을 갖추고 있습니다. 그는 Llama 모델 지원과 RoPE 스케일링, Speculative Decoding과 같은 기능을 구현하며 TensorRT-LLM 개발에도 기여했습니다. Rakib은 Louisiana State University(LSU)에서 x64 및 ARM CPU 기반 수학 라이브러리 최적화를 주제로 박사 학위를 취득했습니다.
Alessio Devoto 프로필
Alessio Devoto는 Sapienza University of Rome에서 데이터 사이언스 분야 박사과정 연구원으로, 효율적이고 해석 가능한 머신러닝 모델 개발에 집중하고 있습니다. 그의 연구는 컴퓨터 비전과 자연어 처리 전반에 걸쳐 있으며, 효율적인 언어·비전 모델, AI 해석 가능성, 적응형 연산 기법에 초점을 둡니다. 이전에 NVIDIA 뮌헨 오피스에서 Applied Agent Research 인턴으로 근무했으며, Edinburgh NLP에서 방문 연구원으로 활동한 바 있습니다.
Darragh Hanley 프로필
Darragh Hanley는 NVIDIA Applied Agent Research 팀의 Senior LLM Technologist로, 자율적 추론 및 멀티 에이전트 협업 프레임워크 개발을 담당하고 있습니다. 그는 여러 대회 우승 경험을 보유한 NVIDIA Kaggle Grandmaster 팀의 일원이며(최고 순위: 전 세계 4위), UC Berkeley 데이터사이언스 석사 프로그램에서 딥러닝을 가르친 바 있습니다.

댓글