dustynv/mlc

mlc

镜像概述和主要用途

该容器为https://github.com/mlc-ai/mlc-llm%E9%A1%B9%E7%9B%AE%E6%8F%90%E4%BE%9B%E8%BF%90%E8%A1%8C%E7%8E%AF%E5%A2%83%EF%BC%8C%E5%9F%BA%E4%BA%8EApache TVM Unity构建，集成了CUDA、cuDNN、CUTLASS、FasterTransformer和FlashAttention-2等加速内核，适用于在支持L4T系统上进行大语言模型的量化处理和性能评估。

核心功能和特性

• 支持大语言模型的W4A16量化（如q4f16_ft量化方案）
• 集成多种GPU加速技术，优化模型推理性能
• 提供模型性能基准测试工具，可评估预填充时间和生成速率
• 兼容Hugging Face模型生态，支持主流LLM模型（如Llama-2、Llava等）

使用场景和适用范围

• 大语言模型在边缘设备（如Jetson系列）上的部署前量化优化
• 评估不同量化配置下模型的推理性能
• 为MLC LLM项目开发提供一致的运行环境

使用方法和配置说明

模型量化

步骤1：准备原始模型

首先下载需要量化的Hugging Face Transformers格式模型，并创建符号链接至/data/models/mlc/dist/models目录，使MLC能正确识别：

bash
./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'

[!NOTE]
若量化Llava模型，需将原始模型的config.json中"model_type": "llava"修改为"model_type": "llama"（可在模型下载至/data/models/huggingface后本地修改）

步骤2：执行量化

对模型执行W4A16量化：

bash
./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

量化后的模型及其运行时文件将保存至/data/models/mlc/dist/Llama-2-7b-chat-hf-q4f16_ft目录。

性能基准测试

使用benchmark.py脚本评估量化模型性能：

bash
./run.sh $(./autotag mlc) \
  python3 /opt/mlc-llm/benchmark.py \
    --model /data/models/mlc/dist/Llama-2-7b-chat-hf-q4f16_ft/params \
    --prompt /data/prompts/completion_16.json \
    --max-new-tokens 128

参数说明

• --prompt：指定输入提示文件，控制上下文长度（/data/prompts目录下提供最长4096 tokens的提示序列）
• --max-new-tokens：指定模型为每个提示生成的输出tokens数量

输出示例

AVERAGE OVER 10 RUNS:
/data/models/mlc/dist/Llama-2-7b-chat-hf-q4f16_ft/params:  prefill_time 0.027 sec, prefill_rate 582.8 tokens/sec, decode_time 2.986 sec, decode_rate 42.9 tokens/sec

结果解释

• 预填充时间（prefill_time）：模型处理输入上下文并开始生成输出前的耗时
• 生成速率（decode_rate）：模型生成输出tokens的速度
• 结果为多次运行的平均值（排除首次运行作为热身）

容器信息

容器镜像列表

运行容器

bash
# 自动拉取或构建兼容的容器镜像
jetson-containers run $(autotag mlc)

# 或显式指定上述镜像之一
jetson-containers run dustynv/mlc:0.1.1-r36.2.0

# 或使用'docker run'（需指定镜像及挂载等参数）
sudo docker run --runtime nvidia -it --rm --network=host dustynv/mlc:0.1.1-r36.2.0

bash
jetson-containers run -v /host/path:/container/path $(autotag mlc)

bash
jetson-containers run $(autotag mlc) my_app --abc xyz

构建容器

bash
jetson-containers build mlc