Turbinando o Desempenho do Meta Llama 3 com NVIDIA TensorRT-LLM e Servidor de Inferência NVIDIA Triton

por Anjali Shah

Temos o prazer de anunciar o suporte para a família de modelos Meta Llama 3 no NVIDIA TensorRT-LLM, acelerando e otimizando seu desempenho de inferência LLM. Você pode experimentar imediatamente o Llama 3 8B e o Llama 3 70B (os primeiros modelos da série) por meio de uma interface de usuário do navegador. Ou, por meio de endpoint de API executados em um stack NVIDIA totalmente acelerado do catálogo de API da NVIDIA, onde o Llama 3 é empacotado como um NVIDIA NIM com uma API padrão que pode ser implantada em qualquer lugar.

Grandes modelos de linguagem são computacionalmente intensivos. Seu tamanho os torna caros e lentos para serem executados, especialmente sem as técnicas certas. Muitas técnicas de otimização estão disponíveis, como kernel fusion e quantização para otimizações de tempo de execução, como implementações C++, cache KV, envio contínuo em lote em voo e atenção paginada. Os desenvolvedores devem decidir qual combinação ajuda seu caso de uso. O TensorRT-LLM simplifica esse trabalho.

O TensorRT-LLM é uma biblioteca de código aberto que acelera o desempenho de inferência nos LLMs mais recentes em GPUs NVIDIA. O NeMo, um framework de ponta a ponta para criar, personalizar e implantar aplicações de IA generativa, usa o TensorRT-LLM e o Servidor de Inferência NVIDIA Triton para implantações de IA generativa.

O TensorRT-LLM usa o compilador de deep learning NVIDIA TensorRT. Ele inclui os mais recentes kernels otimizados para implementações de ponta do FlashAttention e masked multi-head attention (MHA) para execução do modelo LLM. Ele também consiste em etapas de pré e pós-processamento e primitivas de comunicação multi-GPU/multi-nó em uma API Python simples e de código aberto para desempenho inovador de inferência LLM em GPUs.

Para ter uma ideia da biblioteca e como usá-la, vamos ver um exemplo de como usar e implantar o Llama 3 8B com o TensorRT-LLM e o Servidor de Inferência Triton.

Para ter uma visão mais detalhada, incluindo diferentes modelos, diferentes otimizações e execução de várias GPUs, confira a lista completa de exemplos do TensorRT-LLM.

Introdução à Instalação

Começaremos clonando e criando a biblioteca TensorRT-LLM seguindo as instruções de instalação específicas do sistema operacional com o comando pip. Esta é uma das maneiras mais fáceis de construir o TensorRT-LLM. Como alternativa, a biblioteca pode ser instalada usando um dockerfile para recuperar dependências.

Os comandos a seguir extraem a biblioteca de código aberto e instalam as dependências necessárias para instalar o TensorRT-LLM dentro do contêiner.

git clone -b v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM

Recuperando os Pesos do Modelo

TensorRT-LLM é uma biblioteca para inferência de LLM. Para usá-lo, você deve fornecer um conjunto de pesos treinados. Um conjunto de pesos pode ser extraído de repositórios como o Hugging Face Hub ou o NVIDIA NGC. Outra opção é usar seus próprios pesos de modelo treinados em um framework como o NeMo.

Os comandos neste post puxam automaticamente os pesos (e arquivos tokenizer) para a variante ajustada de instruções do modelo Llama 3 de 8 bilhões de parâmetros do Hugging Face Hub. Você também pode fazer o download dos pesos para usar offline com o seguinte comando e atualizar os caminhos em comandos posteriores para apontar para este diretório:

git lfs install
git clone https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct

Observe que o uso desse modelo está sujeito a uma licença específica. Concorde com os termos e autentique-se com o HuggingFace para fazer o download dos arquivos necessários.

Executando o Contêiner TensorRT-LLM

Iniciaremos um contêiner docker base e instalaremos as dependências exigidas pelo TensorRT-LLM.

# Obtain and start the basic docker image environment.
docker run --rm --runtime=nvidia --gpus all --volume ${PWD}:/TensorRT-LLM --entrypoint /bin/bash -it --workdir /TensorRT-LLM nvidia/cuda:12.1.0-devel-ubuntu22.04

# Install dependencies, TensorRT-LLM requires Python 3.10
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev

# Install the stable version (corresponding to the cloned branch) of TensorRT-LLM.
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com

Compilando o Modelo

A próxima etapa no processo é compilar o modelo em um mecanismo TensorRT com pesos de modelo e uma definição de modelo escrita na API Python TensorRT-LLM.

O repositório TensorRT-LLM inclui várias arquiteturas de modelo e usamos a definição de modelo Llama. Para obter mais detalhes e plug-ins e quantizações mais robustos disponíveis, consulte este exemplo de Llama e a documentação de precisão.

# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface-cli login --token *****

# Build the Llama 8B model using a single GPU and BF16.
python3 examples/llama/convert_checkpoint.py --model_dir ./Meta-Llama-3-8B-Instruct \
            --output_dir ./tllm_checkpoint_1gpu_bf16 \
            --dtype bfloat16

trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_bf16 \
            --output_dir ./tmp/llama/8B/trt_engines/bf16/1-gpu \
            --gpt_attention_plugin bfloat16 \
            --gemm_plugin bfloat16

Quando criamos a definição do modelo com a API TensorRT-LLM, construímos um gráfico de operações a partir de primitivas TensorRT que formam as camadas de nossa rede neural. Essas operações são mapeadas para kernels específicos que são programas pré-escritos para a GPU.

O compilador TensorRT pode varrer o gráfico para escolher o melhor kernel para cada operação e cada GPU disponível. Ele também pode identificar padrões no gráfico onde várias operações são boas candidatas para fusão em um único kernel fundido, reduzindo a quantidade necessária de movimento de memória e a sobrecarga de lançamento de vários kernels de GPU.

Além disso, o TensorRT constrói o gráfico de operações em um NVIDIA CUDA Graph que pode ser iniciado ao mesmo tempo. Isso reduz ainda mais a sobrecarga de lançamento de kernels.

O compilador TensorRT é eficiente em fundir camadas e aumentar a velocidade de execução, no entanto, existem algumas fusões de camadas complexas, como o FlashAttention, que envolvem a intercalação de muitas operações e que não podem ser descobertas automaticamente. Para esses, podemos substituir explicitamente partes do gráfico por plug-ins em tempo de compilação. Em nosso exemplo, incluímos o plug-in gpt_attention, que implementa um kernel de atenção fundida semelhante ao FlashAttention, e o plug-in gemm, que executa a multiplicação de matriz com acúmulo de FP32. Também destacamos nossa precisão desejada para o modelo completo como FP16, correspondendo à precisão padrão dos pesos que fazemos o download do HuggingFace.

Quando terminarmos de executar o script de compilação, devemos esperar ver os três arquivos a seguir na pasta /tmp/llama/8B/trt_engines/bf16/1-gpu:

  • rank0.engine é a saída principal do nosso script de construção, contendo o gráfico executável de operações com os pesos do modelo incorporados.
  • config.json inclui informações detalhadas sobre o modelo, como sua estrutura geral e precisão, bem como informações sobre quais plug-ins foram incorporados ao motor.

Executando o Modelo

Então, agora que temos nosso motor modelo, o que podemos fazer com ele?

O arquivo de mecanismo contém as informações para executar o modelo. O TensorRT-LLM inclui um tempo de execução C++ altamente otimizado para executar arquivos de mecanismo e gerenciar processos como tokens de amostragem da saída do modelo, gerenciamento do cache KV e solicitações em lote juntos.

Podemos usar o tempo de execução diretamente para executar o modelo localmente ou podemos implantar usando o Servidor de Inferência Triton no ambiente de produção para compartilhar o modelo com vários usuários.

Para executar o modelo localmente, podemos executar o seguinte comando:

python3 examples/run.py --engine_dir=./tmp/llama/8B/trt_engines/bf16/1-gpu --max_output_len 100 --tokenizer_dir ./Meta-Llama-3-8B-Instruct --input_text "How do I count to nine in French?"

Implantando com o Servidor de Inferência Triton

Além da execução local, também podemos usar o Servidor de Inferência Triton para criar uma implantação pronta para produção de nosso LLM. O back-end do Servidor de Inferência Triton para TensorRT-LLM usa o tempo de execução do TensorRT-LLM C++ para execução de inferência de alto desempenho. Ele inclui técnicas como batching em voo e cache KV paginado que fornecem alta taxa de transferência com baixa latência. O back-end do TensorRT-LLM foi empacotado com o Servidor de Inferência Triton e está disponível como um contêiner pré-construído no NGC.

Primeiro, devemos criar um repositório de modelo para que o Servidor de Inferência Triton possa ler o modelo e quaisquer metadados associados.

O repositório tensorrtllm_backend inclui a configuração de um repositório de modelo necessário em all_models/inflight_batcher_llm/ que podemos replicar.

No diretório há quatro subpastas contendo artefatos para diferentes partes do processo de execução do modelo. As pastas de preprocessing/ e postprocessing/ contêm scripts para o back-end python do Servidor de Inferência Triton. Esses scripts servem para tokenizar as entradas de texto e destokenizar as saídas do modelo para converter entre cadeias de caracteres e as IDs de token nas quais o modelo opera.

A pastatensorrt_llm é onde colocaremos o mecanismo de modelo que compilamos anteriormente. E, finalmente, a pasta  ensemble define um conjunto de modelos que vincula os três componentes anteriores e diz ao Servidor de Inferência Triton como fluir dados através deles.

Puxe para baixo o repositório de modelo de exemplo e copie o modelo que você compilou na etapa anterior para ele.

# After exiting the TensorRT-LLM docker container
cd ..
git clone -b v0.8.0 https://github.com/triton-inference-server/tensorrtllm_backend.git
cd tensorrtllm_backend
cp ../TensorRT-LLM/tmp/llama/8B/trt_engines/bf16/1-gpu/* all_models/inflight_batcher_llm/tensorrt_llm/1/

Em seguida, devemos modificar os arquivos de configuração do esqueleto do repositório com a localização do mecanismo de modelo compilado. Também devemos atualizar parâmetros de configuração, como tokenizador, para usar e manipular a alocação de memória para o cache KV ao enviar solicitações de inferência em lote.

#Set the tokenizer_dir and engine_dir paths
HF_LLAMA_MODEL=TensorRT-LLM/Meta-Llama-3-8B-Instruct
ENGINE_PATH=tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/preprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,preprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/postprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,postprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/ensemble/config.pbtxt triton_max_batch_size:64

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0

Agora, podemos girar nosso contêiner docker e iniciar o servidor Triton. Devemos especificar o tamanho do mundo (o número de GPUs para as quais o modelo foi construído) e apontar para nossos model_repo que acabamos de configurar.

#Change to base working directory
cd..
docker run -it --rm --gpus all --network host --shm-size=1g \
-v $(pwd):/workspace \
--workdir /workspace \
nvcr.io/nvidia/tritonserver:24.03-trtllm-python-py3

# Log in to huggingface-cli to get tokenizer
huggingface-cli login --token *****

# Install python dependencies
pip install sentencepiece protobuf

# Launch Server

python3 tensorrtllm_backend/scripts/launch_triton_server.py --model_repo tensorrtllm_backend/all_models/inflight_batcher_llm --world_size 1

Envio de Solicitações

Para enviar solicitações de inferência e receber conclusões do servidor em execução, você pode usar uma das bibliotecas de cliente do Servidor de Inferência Triton ou enviar solicitações HTTP para o endpoint gerado.

O comando curl a seguir demonstra um teste rápido para solicitar conclusões pelo servidor em execução e o script de cliente com mais recursos pode ser revisado para comunicação com o servidor.

curl -X POST localhost:8000/v2/models/ensemble/generate -d \
'{
"text_input": "How do I count to nine in French?",
"parameters": {
"max_tokens": 100,
"bad_words":[""],
"stop_words":[""]
}
}'

Conclusão

O TensorRT-LLM fornece ferramentas para otimizar e executar com eficiência grandes modelos de linguagem em GPUs NVIDIA. O Servidor de Inferência Triton é ideal para implantar e atender com eficiência grandes modelos de linguagem, como o Llama 3.

Use este guia de introdução e comece sua jornada aproveitando ferramentas de código aberto para usar o Llama 3 e muitos outros grandes modelos de linguagem.

O NVIDIA AI Enterprise, uma plataforma de software de IA de ponta a ponta que inclui o TensorRT, em breve incluirá o TensorRT-LLM, para inferência de IA de missão crítica com segurança, estabilidade, capacidade de gerenciamento e suporte de nível empresarial.