大型語言模型 (LLM) 的出現革新了人工智能領域，為與數字世界的交互提供了全新的方式。盡管 LLM 通常能夠提供良好的通用解決方案，但為了更好地支持特定領域和任務，它們往往需要進行調整。
AI 編碼助手（或代碼 LLM）已成為幫助實現這一目標的一個領域。到 2025 年，80% 的產品開發生命周期將使用 生成式 AI 進行代碼生成，開發者將充當后端和前端組件及集成的驗證者和編排者。您可以調整用于代碼任務的 LLM，簡化開發者的工作流程，并降低新手編程人員的門檻。Code LLM 不僅可以生成代碼，還可以填充缺失的代碼、添加文檔，并提供解決難題的提示。
本文將介紹如何部署端到端代碼 LLM，包括具體的提示指南、優化技術和客戶端 – 服務器部署。我們使用NVIDIA Triton 推理服務器并使用NVIDIA TensorRT-LLM，這是一個用于編譯和優化用于推理的 LLM 的綜合庫。TensorRT-LLM 包含許多高級優化，同時提供直觀的 Python API 來定義和構建新模型。
我們TensorRT-LLM 開源庫加速 NVIDIA GPU 上最新 LLM 的推理性能。它用作 NVIDIA GPU 中 LLM 推理的優化主干NVIDIA NeMo，這是一種端到端框架，用于構建、自定義生成式 AI 應用并將其部署到生產環境中。NeMo 框架為生成式 AI 部署提供完整的容器，包括 TensorRT-LLM 和 NVIDIA Triton 推理服務器。它還包含在NVIDIA AI Enterprise中，這是一個生產級、安全且受支持的軟件平臺，可簡化采用流程。
什么是 AI 輔助編碼
程序員需要花費大量時間來尋找解決方法來解決常見問題，或瀏覽在線論壇，以更快地編寫代碼。AI 代碼助手背后的關鍵理念是將程序員所需的信息放在他們正在編寫的代碼旁邊。該工具會跟蹤程序員正在處理的文件中的代碼和評論，以及它鏈接到或在同一項目中已編輯的其他文件。它會將所有這些文本作為提示發送到 LLM。然后，它會預測程序員正在嘗試執行的操作，并建議代碼來完成此操作。
部署自己的 AI 編碼助手
要部署自己的 AI 編碼助手，您需要深度學習推理和 LLM 的基本知識，以及：
- 擁抱 Face 注冊用戶訪問權限并基本熟悉 Transformer 庫
- Python
- NVIDIA TensorRT-LLM 優化庫
- 搭載 TensorRT-LLM 后端的 NVIDIA Triton
本教程使用 StarCoder，這是一個 155 億個參數 LLM，使用 The Stack (v1。2)中的 80 多種編程語言進行訓練。StarCoder 的基礎模型使用來自 80 多種編程語言、GitHub 問題、Git Commits 和 Jupyter Notebooks 的 1 萬億個令牌進行訓練。StarCoder 在其之上使用另外 350 億個 Python 令牌進行了微調，最終形成了 StarCoder 模型。
針對用于解決特定問題的代碼請求，StarCoder 可以從頭開始生成代碼。它還支持中間填充 (FIM) 或填充，這使得模型能夠根據插入點周圍的前綴和后綴代碼生成代碼(例如，單行填充、返回類型預測和文檔字符串生成)。
代碼 LLM 提示
本節介紹優化過程。StarCoder 在公共領域的代碼庫上進行了訓練，因此可以生成多種語言的代碼片段。LLM 的基于代碼的提示是要求 LLM 以特定語言生成代碼的提示。理想的提示應該具體清晰，為 LLM 提供足夠的信息以生成正確答案。
例如，您可以發出以下通用提示：
Write a function that computes the square root。
# Use Newton's method,# where x_(n+1) = 1/2 * (x_n + (y/x_n))# y = number to find square root of# x_0 = first guess# epsilon = how close is close enough?# Implement this in a function called newton_sqrt that has three parameters# and returns one value. |
指定編程語言和輸出將產生更好的結果。例如：
Write Python code to compute the square root and print the result。
# To find the square root of a number in Python, you can use the math library and its sqrt function:?
from math import sqrt?
number = float(input('Enter a number: '))square_root = sqrt(number)print(f'The square root of {number} is approximately {square_root:.2f}.') |
代碼 LLM 的提示工程
提示工程包括提供代碼塊、單行代碼或開發者評論，以改進特定用例的編碼建議。在針對代碼 LLM 的特定提示工程方面，有一些指南需要記住，包括代碼片段和示例輸出。
片段
通常，LLM 的上下文窗口大小有限。這對于代碼 LLM 尤其重要，因為源代碼通常由多個文件和長文本組成，通常不適合 LLM 上下文窗口。因此，代碼 LLM 的提示工程需要在提示之前為源代碼做額外的準備。這可以包括將文件拆分為較小的代碼片段、選擇最具代表性的代碼示例，以及使用提示優化技術進一步減少輸入令牌的數量。
添加示例輸出
包含示例輸出和導入指令的提示將產生最佳結果。
您將通過以下提示獲得平均結果：
Write a function for calculating average water use per household。
添加示例將產生更好的結果：
Write a function for calculating average water use per household。Example output ["Smith", 20, "Lincoln", 30, "Simpson", 1500]
通過添加示例并指定導入說明，您的結果將更加出色：
Write a function for calculating average water use per household。Example input [['Last Name', 'Dishwasher', 'Shower', 'Yard'], ['Smith', 39, 52, 5], ['Lincoln', 25, 77, 8], ['Simpson', 28, 20, 0]]Example output ["Smith", 20, "Lincoln", 30, "Simpson", 1500]
實驗
使用不同的提示，并在每個步驟中使用更精細的細節進行更新，以提高輸出。例如：
Write a function for calculating average water use per household。Write a function for calculating average water use per household. Add penalties for dishwasher time。Write a Python function for calculating average water use per household. Add penalties for dishwasher time. Input is family name, address, and number of gallons per specific use。
流程概述
要使用 TensorRT-LLM 優化 LLM，您需要了解其架構，并確定它與哪個通用基礎架構最相似。StarCoder 使用 GPT 架構，因此本教程基于 NVIDIA/TensorRT-LLM GPT 示例。
使用此目錄中的轉換器和構建腳本編譯 StarCoder 并為硬件加速做好準備。請注意，分詞器并非由 TensorRT-LLM 直接處理。但是，必須能夠在定義的分詞器系列中對其進行分類，以用于運行時以及在 Triton 中設置預處理和后處理步驟。
設置和構建 TensorRT-LLM
首先，克隆和構建 TensorRT-LLM 庫。構建 TensorRT – LLM 并檢索其所有依賴項的最簡單方法是使用隨附的 Dockerfile。這些命令拉取基礎容器，并在容器內安裝 TensorRT – LLM 所需的所有依賴項。然后，它會在容器內構建并安裝 TensorRT – LLM 本身。
git lfs installgit clone -b release/0.7.1 https://github.com/NVIDIA/TensorRT-LLM.gitcd TensorRT-LLMgit submodule update --init --recursivemake -C docker release_build |
檢索模型權重
從 Hugging Face 下載 StarCoder 模型，并將其放置在/examples目錄。
cd examplesgit clone https://huggingface.co/bigcode/starcoder |
轉換模型權重格式
在編譯之前，需要將模型權重格式從 Hugging Face Transformer 轉換為 FasterTransformer。
# Launch the Tensorrt-LLM containermake -C docker release_run LOCAL_USER=1?
cd examples/gpt?
python3 hf_gpt_convert.py -p 8 --model starcoder -i ../starcoder -o ./c-model/starcoder --tensor-parallelism 1 --storage-type float16 |
編譯模型
下一步是使用上一步中轉換的權重將模型編譯為 TensorRT 引擎。編譯有許多選項，例如精度和選擇要啟用的優化功能(例如動態批處理和 KV 緩存)。
python3 examples/gpt/build.py \????--model_dir examples/gpt/c-model/starcoder/1-gpu \????--dtype float16 \????--use_gpt_attention_plugin float16 \????--use_gemm_plugin float16 \????--remove_input_padding \????--use_inflight_batching \????--paged_kv_cache \????--output_dir examples/gpt/out |
運行模型
TensorRT-LLM 包含高度優化的 C++運行時，用于執行構建的 LLM 引擎和管理流程，例如從模型輸出中采樣令牌、管理 KV 緩存以及同時處理批處理請求。您可以直接使用該運行時在本地執行模型。
python3 examples/gpt/run.py --engine_dir=examples/gpt/out --max_output_len 100 --tokenizer examples/starcoder/starcoder --input_text "Write a function that computes the square root." |
您可以使用更復雜的提示進行實驗，包括函數名稱、參數和輸出的定義。
python3 examples/gpt/run.py --engine_dir=examples/gpt/out --max_output_len 100 --tokenizer examples/starcoder/starcoder --input_text "X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.1)# Train a logistic regression model, predict the labels on the test set and compute the accuracy score" |
在 NVIDIA Triton 上部署模型
我們建議使用 NVIDIA Triton 推理服務器，這是一個開源平臺，可簡化和加速 AI 推理工作負載的部署，以創建 LLM 的生產就緒型部署。這將有助于縮短設置和部署時間。適用于 TensorRT-LLM 的 Triton 推理服務器后端利用 TensorRT-LLM C++運行時實現快速推理執行，并包含 動態批處理 和 分頁 KV 緩存。您可以通過 NVIDIA NGC 目錄中訪問 TensorRT-LLM 后端用作預構建容器的 Triton 推理服務器 。
首先，創建一個模型庫，以便Triton可以讀取模型和任何相關元數據。tensorrtllm_backend 存儲庫包含 all_models/inflight_batcher_llm/ 下適當模型存儲庫的骨架。該目錄中有以下子文件夾，其中包含模型執行過程不同部分的構件：
/preprocessing和/postprocessing：包含適用于 Python 的 Triton 后端，用于在字符串和模型運行所用的標記 ID 之間進行轉換，實現文本輸入的標記化和模型輸出的去標記化。/tensorrt_llm：用于存儲您之前編譯的模型引擎的子文件夾。/ensemble：定義模型集成，它將前三個組件連接在一起，并告訴 Triton 如何通過它們流動數據。
# After exiting the TensorRT-LLM docker containercd ..git clone -b release/0.7.1 \https://github.com/triton-inference-server/tensorrtllm_backend.gitcd tensorrtllm_backendcp ../TensorRT-LLM/examples/gpt/out/*?? all_models/inflight_batcher_llm/tensorrt_llm/1/ |
密切關注參數 kv_cache_free_gpu_mem_fraction，該參數用于在推理期間為 KV 緩存分配顯存。此參數設置的數字介于 0 和 1 之間，并相應地分配 GPU 顯存。將其設置為 1 將使 GPU 顯存容量達到最大值。
在部署大型 LLM 或在單個 GPU 上運行多個 TensorRT-LLM 實例時，設置此參數尤為重要。請注意，默認值為 0.85.對于本教程中的 StarCoder，我們使用了 0。2 的值以避免 GPU 顯存填充。
借助 KV 緩存，中間數據會被緩存和重復使用，從而避免重新計算。中間特征圖被緩存并在下一次迭代中重復使用，而不是在每次迭代中重新計算完整的鍵值矩陣。這降低了解碼機制的計算復雜性。KV 緩存的內存消耗可以輕松填充 TB 級數據，以用于生產級 LLM。
python3 tools/fill_template.py --in_place \??????all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt \??????decoupled_mode:true,engine_dir:/all_models/inflight_batcher_llm/tensorrt_llm/1,\max_tokens_in_paged_kv_cache:,batch_scheduler_policy:guaranteed_completion,kv_cache_free_gpu_mem_fraction:0.2,\max_num_sequences:4 |
另一個需要注意的重要參數是分詞器。您需要指定合適的分詞器及其最適合的類別類型，以定義 Triton 所需的預處理和后處理步驟。StarCoder 使用代碼作為輸入，而非句子，因此它具有唯一的分詞和分隔符。
因此，有必要在配置文件中通過設置 StarCoder 分詞器tokenizer_type?為自動和tokenizer_dir為/all_models/startcoder 。您可以通過調用fill_template.py或編輯config.pbtxt文件。
python tools/fill_template.py --in_place \????all_models/inflight_batcher_llm/preprocessing/config.pbtxt \????tokenizer_type:auto,tokenizer_dir:/all_models/startcoder??python tools/fill_template.py --in_place \????all_models/inflight_batcher_llm/postprocessing/config.pbtxt \????tokenizer_type:auto,tokenizer_dir:/all_models/startcoder |
啟動 Triton
如要啟動 Triton，請先啟動 Docker 容器，然后通過設置world_size參數。
docker run -it --rm --gpus all --network host --shm-size=1g \-v $(pwd)/all_models:/all_models \-v $(pwd)/scripts:/opt/scripts \nvcr.io/nvidia/tritonserver:23.10-trtllm-python-py3??# Log in to huggingface-cli to get tokenizerhuggingface-cli login --token *****??# Install python dependenciespip install sentencepiece protobuf??# Launch Serverpython /opt/scripts/launch_triton_server.py --model_repo /all_models/inflight_batcher_llm --world_size 1 |
發送請求
要向正在運行的服務器發送請求并與之交互，您可以使用 Triton 客戶端庫 或將 HTTP 請求發送到 生成的端點。
要開始處理簡單請求，請使用以下 curl 命令將 HTTP 請求發送到生成的端點：
curl -X POST localhost:8000/v2/models/ensemble/generate -d \'{"text_input": "write in python code that plots in a image circles with different radiuses","parameters": {"max_tokens": 100,"bad_words":[""],"stop_words":[""]}}' |
您還可以使用功能齊全的客戶端腳本與正在運行的服務器進行交互。
運行上述命令將生成以下代碼。生成的圖形如圖 1 所示。
import numpy as npimport matplotlib.pyplot as plt?
# Fixing random state for reproducibilitynp.random.seed(19680801)?
N = 100r0 = 0.6x = np.random.rand(N)y = np.random.rand(N)area = np.pi * (10 * np.random.rand(N))**2 # 0 to 10 point radiusesc = np.random.rand(N)?
plt.scatter(x, y, s=area, c=c, alpha=0.5)plt.show() |
開始使用
NVIDIA TensorRT-LLM 和 NVIDIA Triton 為眾多流行的 LLM 架構提供了基準支持，使得部署、優化和運行各種代碼 LLM 變得輕松。要開始使用，請下載并設置 NVIDIA/TensorRT-LLM GitHub 上的開源庫，并嘗試使用不同的 LLM 示例。此外，您還可以下載 StarCoder 模型，并按照此博文中的步驟進行動手實踐。
?
