Inferentia에서 Triton 사용하기

Inferentia에서 Triton 사용하기 (Using Triton with Inferentia)

21.11 릴리스부터 트리톤은 AWS Inferentia와 Neuron Runtime을 지원해요. AWS의 인퍼런스 전용 칩인 Inferentia에서 트리톤을 돌리는 전체 흐름을 살펴보죠.

Inferentia 설정

첫 단계는 Deep Learning AMI로 AWS Inferentia 인스턴스를 만드는 거예요(Ubuntu 18.04로 테스트).

ssh   -i   <private-key-name>.pem   ubuntu@<instance   address>

[!참고] 저장 공간을 기본값 110 GiB보다 크게 잡는 걸 권장해요. 현재 버전의 트리톤은 500 GiB 저장소로 테스트됐습니다.

inf1* 인스턴스에 로그인한 뒤 이 currrent GitHub 리포를 클론하고, /home/ubuntu에 둘게요.

chmod 777 /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh
sudo /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh

그다음 트리톤 인스턴스를 시작합니다.

docker run --device /dev/neuron0 <more neuron devices> -v /home/ubuntu/python_backend:/home/ubuntu/python_backend -v /lib/udev:/mylib/udev --shm-size=1g --ulimit memlock=-1 -p 8000:8000 -p 8001:8001 -p 8002:8002 --ulimit stack=67108864 -ti nvcr.io/nvidia/tritonserver:<xx.yy>-py3

• 주의 1: 컨테이너 초기화 시 사용할 neuron 장치를 각각 나열해야 해요. 예를 들어 인스턴스에서 neuron 장치 4개를 쓰려면 --device /dev/neuron0 --device /dev/neuron1 --device /dev/neuron2 --device /dev/neuron3 ...처럼 지정합니다.
• 주의 2: /mylib/udev는 Neuron 파라미터 전달에 사용됩니다.
• 주의 3: 트리톤 컨테이너 버전 <xx.yy>는 Triton Inference Server Container Release Notes를 참고하세요. 현재 빌드 스크립트는 컨테이너 버전 21.10으로 테스트됐어요.

트리톤 컨테이너를 시작한 뒤 python_backend 폴더로 들어가 설정 스크립트를 실행합니다.

source /home/ubuntu/python_backend/inferentia/scripts/setup.sh

이 스크립트는:

• 필요한 의존성 설치
• Neuron 컴파일러인 neuron-cc 설치
• pytorch, tensorflow 또는 둘 다 중 선택한 대로 neuron 프레임워크 패키지 설치

-h 또는 --help 옵션으로 더 많은 구성 옵션을 확인할 수 있어요.

Inferentia 모델 설정

현재 inferentia 실행은 PyTorch와 TensorFlow 워크플로만 지원합니다. 사용자는 자신의 *.pt(pytorch) 또는 *.savedmodels(tensorflow) 모델을 만들어야 해요. Inferentia는 추론 요청 실행에 내부 .NEFF 그래프가 필요하므로 이 단계가 중요합니다. 모델 컴파일은 Neuron 컴파일러 CLI 레퍼런스, PyTorch-Neuron trace 파이썬 API, PyTorch·TensorFlow 튜토리얼을 참고하세요.

PyTorch

Inferentia 실행을 위해 PyTorch-Neuron trace 파이썬 API로 트레이스된 모델을 지원합니다. Inferentia를 지원하는 TorchScript 모델을 얻으면 gen_triton_model.py 스크립트로 트리톤 파이썬 모델 디렉토리를 생성해요.

PyTorch 모델용 예시 호출:

python3 inferentia/scripts/gen_triton_model.py --model_type pytorch --triton_input INPUT__0,INT64,4x384 INPUT__1,INT64,4x384 INPUT__2,INT64,4x384 --triton_output OUTPUT__0,INT64,4x384 OUTPUT__1,INT64,4x384 --compiled_model /home/ubuntu/bert_large_mlperf_neuron_hack_bs1_dynamic.pt --neuron_core_range 0:3 --triton_model_dir bert-large-mlperf-bs1x4

컴파일된 모델을 TorchScript 모델로 취급하려면 --model_type pytorch를 제공해야 해요.

[!참고] TorchScript 모델에는 입출력 메타데이터(이름·데이터타입·텐서 shape)가 없어서 위 스크립트에 모두 제공해야 하고, 이름은 <name>__<index> 컨벤션을 따라야 해요. <name>은 아무 문자열, <index>는 입출력의 위치입니다. 입출력이 둘씩이면 INPUT__0, INPUT__1, OUTPUT__0, OUTPUT__1처럼 이름을 지어요.

추가로 --neuron_core_range는 이 모델 서빙에 사용할 neuron 코어를 지정합니다. 현재는 torch.neuron.DataParallel() 모드만 지원돼요. 트리톤 모델 인스턴스 수는 --triton_model_instance_count 옵션으로 지정하고, neuron 코어는 모든 인스턴스에 균등 분배됩니다. 예를 들어 트리톤 인스턴스 2개에 neuron 코어 4개면 첫 인스턴스는 코어 0-1, 둘째는 2-3에 로딩돼요. Inferentia 장치를 가장 잘 활용하려면 neuron 코어 수를 인스턴스 수의 배수로 잡는 게 좋습니다.

TensorFlow

TensorFlow 모델은 AWS Neuron용으로 컴파일돼야 해요. AWS Neuron TensorFlow 튜토리얼로 Neuron 코어를 쓰는 컴파일 모델을 만드는 법을 배울 수 있어요. 현재 코드는 tensorflow==1.15에서만 테스트됩니다. 컴파일 모델을 얻은 뒤 gen_triton_model.py로 트리톤 파이썬 모델 디렉토리를 생성합니다.

TensorFlow 모델용 예시 호출:

python3 gen_triton_model.py --model_type tensorflow --compiled_model /home/ubuntu/inferentia-poc-2.0/scripts-rn50-tf-native/resnet50_mlperf_opt_fp16_compiled_b5_nc1/1 --neuron_core_range 0:3  --triton_model_dir rn50-1neuroncores-bs1x1

[!참고] TorchScript 모델과 달리 TensorFlow SavedModel은 입출력 텐서의 이름·데이터타입·shape 메타데이터를 충분히 저장하고 있어요. 기본적으로 스크립트는 컴파일 모델을 torchscript로 간주하므로, TF savedmodel로 취급하려면 --model_type tensorflow를 제공해야 합니다. tensorflow 모델에는 tensorflow 파이썬 모듈 설치가 필요해요.

PyTorch와 유사하게 --neuron_core_range와 --triton_model_instance_count를 쓸 수 있지만, TensorFlow에서 neuron 코어 인덱스는 칩의 특정 코어를 가리키지 않아요. TensorFlow는 폐기된 NEURONCORE_GROUP_SIZES 기능으로 모델을 로딩하며, 이 경우 모델은 다음 가용한 Neuron 코어들에 로딩됩니다. 그리고 Neuron-Tensorflow는 여러 코어용 내장 실행 함수가 없어 model.py가 입력 텐서를 가용 코어에 분할해 워크로드를 분산해요. 여러 코어 처리 시 입력의 첫 번째 차원은 None으로 두는 걸 권장합니다.

gen_triton_model.py의 -h/--help로 더 많은 옵션을 확인하세요.

Inferentia 모델을 트리톤에서 서빙하기

gen_triton_model.py는 다음 구조의 트리톤 모델 디렉토리를 만들어요.

bert-large-mlperf-bs1x4
|
|- 1
|  |- model.py
|
|- config.pbtxt

사용자가 준 이름으로 모델 디렉토리가 생성되므로, 그 디렉토리를 트리톤의 모델 레포지토리로 옮깁니다. 스크립트에 준 컴파일 모델 경로가 유효한 torchscript 파일 또는 tensorflow savedmodel을 가리키는지 확인하세요.

이제 서버를 아래처럼 모델과 함께 실행할 수 있어요.

tritonserver --model-repository <path_to_model_repository>

[!참고]

• config.pbtxt와 model.py는 시작점으로 취급하세요. 필요에 따라 커스터마이즈해도 됩니다.
• Triton Inferentia는 현재 단일 모델로 테스트됐어요.

트리톤 동적 배칭 사용하기

동적 배칭을 활성화하려면 --enable_dynamic_batching 플래그를 지정해야 해요. gen_triton_model.py는 트리톤의 동적 배칭 구성에 다음 세 가지 옵션을 지원합니다.

• --preferred_batch_size: preferred batch size 상세는 모델 구성 문서 참고. 성능 최적화를 위해 참여 neuron 코어 수의 배수로 두는 걸 권장합니다(예: 인스턴스당 neuron 코어 2개면 2·4·6).
• --max_queue_delay_microseconds: 모델 구성 문서 참고.
• --disable_batch_requests_to_neuron: 트리톤이 배칭 요청을 처리하는 비기본 방식을 켭니다. 트리톤 서버 요청이 배칭됐는지와 무관하게 트리톤 백엔드는 각 요청을 neuron에 개별 전송해요. 플래그 없이는 배칭 성능이 좋지 않은 모델에서 성능을 최적화하려 할 때 권장됩니다.

추가로 --max_batch_size는 최대 배칭 한도를 결정합니다. 자세한 내용은 모델 구성 문서를 참고하세요.

Inferentia 설정 정확도 테스트

qa 폴더에는 간단한 add_sub 모델로 테스트를 세팅하는 데 필요한 파일이 있어요. 테스트는 inferentia 코어가 8개넘는 인스턴스(예: inf1.6xlarge)가 필요합니다. 테스트를 시작하려면:

source <triton path>/python_backend/inferentia/qa/setup_test_enviroment_and_test.sh

<triton path>는 보통 /home/ubuntu/입니다. 이 스크립트는 inferentia 테스트가 들어 있는 server 리포를 받고 최신 Triton Server·Triton SDK를 빌드해요.

[!참고] server 리포의 일부 테스트를 바꿔야 한다면 스크립트 실행 전 export TRITON_SERVER_REPO_TAG=<your branch name>을 실행하세요.

Inferentia 2 또는 Trn1

pytorch-neuronx와 tensorflow-neuronx

• inf1과 유사하게, pre-container·on-container 설정 스크립트의 인자에 -inf2 또는 -trn1 플래그를 추가합니다:

chmod 777 /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh
sudo /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh -inf2

• 컨테이너에서 docker run 뒤 setup.sh 스크립트에 비슷한 인자를 전달할 수 있어요.
  • PyTorch: source /home/ubuntu/python_backend/inferentia/scripts/setup.sh -inf2 -p
  • TensorFlow: source /home/ubuntu/python_backend/inferentia/scripts/setup.sh -inf2 -t
• 위 단계 후 gen_triton_model.py를 쓸 때도 --inf2 인자를 전달할 수 있어요(예: PyTorch):

python3 inferentia/scripts/gen_triton_model.py --inf2 --model_type pytorch --triton_input INPUT__0,INT64,4x384 INPUT__1,INT64,4x384 INPUT__2,INT64,4x384 --triton_output OUTPUT__0,INT64,4x384 OUTPUT__1,INT64,4x384 --compiled_model bert_large_mlperf_neuron_hack_bs1_dynamic.pt --neuron_core_range 0:3 --triton_model_dir bert-large-mlperf-bs1x4

[!참고] --inf2 옵션을 쓰면 --compiled_model 경로를 트리톤 모델 디렉토리 기준 상대 경로로 제공해야 해요. model.py의 initialize() 함수가 레포지토리 안 모델 경로와 상대 --compiled_model 경로를 이어 붙여 전체 경로를 유도합니다.

transformers-neuronx

inf2/trn1 인스턴스를 transformers-neuronx 패키지로 서빙하려면 위 지침대로 pytorch 모델을 생성해요. transformers-neuronx는 현재 여기에 나열된 모델들을 지원해요. neuronx 로드 API는 모델마다 다르지만 같은 패턴을 따릅니다.

• transformers-neuronx 모델을 서빙하려면 먼저 inf2 인스턴스(LLM은 inf2.24xl 권장)에서 save_pretrained_split() API로 모델을 트레이스하고, 그 폴더를 gen_triton_model.py의 --compiled_model로 패키징합니다.
• 다음 트리는 OPT 모델의 샘플 구조예요:

opt/
├── 1
│   └── model.py
├── opt-125m-model
│   └── pytorch_model.bin
└── opt-125m-tp12
├── FullyUnrolled.1814.1
│   ├── penguin-sg0000
│   └── sg00
├── FullyUnrolled.1814.2
│   ├── penguin-sg0000
│   └── sg00
├── FullyUnrolled.1814.3
│   ├── penguin-sg0000
│   └── sg00
├── FullyUnrolled.1814.4
│   ├── penguin-sg0000
│   └── sg00
└── FullyUnrolled.1814.5
├── penguin-sg0000
└── sg00
├── config.pbtxt

• 다음 import를 추가해요(OPT 모델 예시, 실행하려는 모델에 따라 다름):

from transformers_neuronx.opt.model import OPTForSampling

• initialize() 함수에 다음 줄을 추가하고 batch_size, tp_degree, n_positions, amp, unroll을 요구사항에 맞게 설정해요. tp_degree는 보통 inf2 인스턴스의 가용 neuron 코어 수와 일치해야 해요.

batch_size = 1
tp_degree = 12
n_positions = 2048
amp = 'bf16'
unroll = None
self.model_neuron = OPTForSampling.from_pretrained(compiled_model, batch_size=batch_size, amp=amp, tp_degree=tp_degree, n_positions=n_positions, unroll=unroll)
self.model_neuron.to_neuron()

self.model_neuron.num_workers = num_threads

batch_size 등의 인자를 config.pbtxt의 파라미터로 넣고 --compiled-model처럼 initialize()에서 읽을 수도 있어요.

• 마지막으로 execute() 함수에서 다음 API로 추론을 실행합니다:

batched_results = self.model_neuron.sample(batched_tensor, 2048)

여기서 2048은 충분히 긴 출력 토큰 수예요. 페이로드의 일부로 지정하고 싶다면 입력 중 하나로 전달할 수도 있습니다.

• 그다음 다른 트리톤 모델처럼 모델을 로딩하고 추론 페이로드를 제출하면 돼요.

더 알아보기 (Learn more)

• Triton Python Backend — 파이썬 백엔드 기반
• AWS Neuron 문서 — 네트워크 컴파일러·프레임워크 가이드
• Triton 모델 구성 — 배칭·인스턴스 그룹 설정