Skip to main content

使用 MLC LLM 在 Jetson 上运行量化的 Llama2-7B

简介​

近年来，大型语言模型（如 GPT-3）在自然语言处理任务中取得了革命性进展。然而，这些模型大多基于大规模数据集进行训练，需要强大的计算资源，难以部署在边缘设备上。为了解决这一问题，研究人员开发了量化技术，将大型模型压缩为更小的模型，同时保持性能不变。

在本项目中，我们介绍了Llama2-7B的量化版本，这是一种基于 1.5TB 数据训练的大型语言模型，并将其部署在 Jetson Orin 上。我们还利用机器学习编译器大型语言模型（MLC LLM）来加速模型的推理速度。通过在 Jetson Orin NX 上部署量化的 Llama2-7B 和 MLC LLM，开发者可以构建高精度、低延迟的边缘设备自然语言处理应用。

硬件组件​

reComputer（或其他基于 Jetson 的设备）
立即获取 🖱️

安装依赖项：​

sudo apt-get update && sudo apt-get install git python3-pip
git clone --depth=1 https://github.com/dusty-nv/jetson-containers
cd jetson-containers pip3 install -r requirements.txt
cd ./data && git clone https://github.com/LJ-Hao/MLC-LLM-on-Jetson-Nano.git && cd ..

安装并运行容器​

第一步：安装镜像​

./run.sh --env HUGGINGFACE_TOKEN=<YOUR-ACCESS-TOKEN> $(./autotag mlc) /bin/bash -c 'ln -s $(huggingface-downloader meta-llama/Llama-2-7b-chat-hf) /data/models/mlc/dist/models/Llama-2-7b-chat-hf'

使用 sudo docker images 检查镜像是否已安装

pir

第二步：安装 Llama2-7b-chat-hf 并使用 MLC 对模型进行量化​

./run.sh $(./autotag mlc) \
python3 -m mlc_llm.build \
--model Llama-2-7b-chat-hf \
--quantization q4f16_ft \
--artifact-path /data/models/mlc/dist \
--max-seq-len 4096 \
--target cuda \
--use-cuda-graph \
--use-flash-attn-mqa

第三步：运行并进入 Docker 容器​

./run.sh <YOUR IMAGE NAME> 
# 对我来说是 dustynv/mlc:51fb0f4-builder-r35.4.1，检查第一步的结果

pir

运行模型​

在未量化的情况下运行 Llama（未使用 MLC LLM 量化）​

cd /data/MLC-LLM-on-Jetson && python3 Llama-2-7b-chat-hf.py

pir

可以看到，在未使用 MLC 量化的情况下，Jetson Nano 16GB 能加载模型但无法运行。

在量化的情况下运行 Llama（使用 MLC LLM 量化）​

cd /data/MLC-LLM-on-Jetson && python3 Llama-2-7b-chat-hf-q4f16_ft.py

运行结果如下：

pir

在 Jetson Orin NX 16GB 上运行 Llama 和 MLC 的视频：​

项目展望​

在本项目中，我们展示了如何在 Jetson Orin 上使用 MLC LLM 部署量化的 Llama2-7B。借助 Jetson Orin 的强大计算能力，开发者可以构建高精度、低延迟的边缘设备自然语言处理应用。未来，我们将继续探索在边缘设备上部署大型语言模型的潜力，并开发更高效、优化的部署方法。

Loading Comments...