TensorRT-LLM在CodeFuse-CodeLlama-34B上的int4量化实践

最新推荐文章于 2025-08-20 17:24:35 发布

原创

最新推荐文章于 2025-08-20 17:24:35 发布 · 1k 阅读

CC 4.0 BY-SA版权

文章标签：

本文详细描述了如何在CodeFuse-CodeLlama-34B模型上进行int4量化，包括TensorRT的安装、GPTQ技术的应用、量化设置以及性能测试，展示了量化后的模型在A10/A100/L40S上的加速效果和内存优化。

本文档是利用TensorRT-LLM(early access)对 CodeFuse-CodeLlama-34B 进行int4量化的实践（注：TensorRT-LLM(GA) 已发布）。Codefuse是由蚂蚁集团开发的专门用于支持整个软件开发生命周期的大型代码语言模型（Code LLMs），涵盖设计、需求、编码、测试、部署、运维等关键阶段。致力于打造创新的解决方案，让软件开发者们在研发的过程中如丝般顺滑。CodeFuse-CodeLlama-34B的是在CodeLlama-34b-Python的基础上，通过高质量指令数据集和多任务范式微调而成的模型，在HumanEval Benchmarks的Python Pass@1 取得了74.4%（greedy decoding）的开源SOTA成绩。下面会展示我们是如何一步一步地完成TensorRT-LLM的安装、量化以及测试。

如果网络能够触达以下url

nvcr.io/nvidia/pytorch:23.07-py3
https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.0.1/tars/TensorRT-9.0.1.4.Linux.x86_64-gnu.cuda-12.2.tar.gz
https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.0.1/tars/polygraphy-0.48.1-py2.py3-none-any.whl

直接使用如下命令进行容器构建

docker build -t tensorrt_llm -f docker/Dockerfile.dev .

但诸如业务生产环境，会出于安全考虑无法直接触达上述url，在这种情况下就需要自己构建镜像。我们是在一个py3.8的pytorch基础镜像上构建tensorrt-llm的运行镜