huggingface/optimum-nvidia

镜像概述

该Docker镜像集成NVIDIA TensorRT-LLM优化引擎与Hugging Face生态，专为大型语言模型（LLM）推理加速设计。通过TensorRT-LLM的高效优化技术（如量化、算子融合、KV缓存优化等），结合Hugging Face Transformers的模型兼容性，提供高性能、低延迟的LLM部署解决方案，适用于需要快速响应的AI服务场景。

核心功能与特性

• TensorRT-LLM优化：集成NVIDIA TensorRT-LLM库，针对LLM进行深度优化，支持INT4/INT8/FP16/FP8等量化精度，显著提升推理吞吐量并降低延迟。

• Hugging Face生态兼容：无缝对接Hugging Face Transformers模型库，支持主流LLM（如Llama、GPT-2、OPT、Falcon等）的加载与部署，无需额外模型转换工具。

• 高性能推理：通过优化的CUDA内核、张量并行与流水线并行支持，充分利用GPU算力，适用于单卡及多卡分布式推理场景。

• 低延迟部署：针对实时推理场景优化，减少模型加载时间与单次推理响应延迟，满足高并发AI服务需求。

• 灵活配置：支持自定义推理参数（如batch size、max sequence length、量化精度等），适配不同业务场景的性能与资源需求。

使用场景与适用范围

• 企业级LLM部署：适用于企业内部知识库问答、智能客服、代码辅助等场景，提供稳定高效的LLM推理服务。

• 实时AI服务：支持需要低延迟响应的应用（如实时聊天机器人、智能内容生成工具），保障用户交互体验。

• 大规模语言模型应用：适配百亿至千亿参数规模的LLM部署，满足科研机构、AI企业的大规模模型验证与服务需求。

使用方法与配置说明

前置条件

• 硬件：NVIDIA GPU（需支持TensorRT-LLM，如A100、H100、L40等），至少16GB显存（根据模型大小调整）。

• 软件：Docker Engine（20.10+）、nvidia-docker runtime、NVIDIA驱动（版本需兼容TensorRT-LLM要求，建议535+）。

基本使用命令

拉取镜像

bash
docker pull [镜像名称]:[标签]  # 请替换为实际镜像名称与标签

启动容器（单模型推理）

bash
docker run -it --rm \
  --gpus all \
  -v /path/to/huggingface/models:/models \  # 挂载Hugging Face模型目录
  -e MODEL_PATH=/models/llama-7b \         # 指定模型路径（容器内路径）
  -e QUANTIZATION=INT8 \                   # 设置量化精度（可选：FP16/INT8/INT4/FP8）
  -e MAX_SEQ_LEN=2048 \                    # 设置最大序列长度
  -p 8080:8080 \                           # 映射推理服务端口
  [镜像名称]:[标签]

环境变量配置

容器支持以下环境变量自定义配置：

Docker Compose示例

创建docker-compose.yml文件，简化多实例或复杂配置部署：

yaml
version: '3.8'
services:
  llm-inference:
    image: [镜像名称]:[标签]
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./local-models:/models  # 本地模型目录挂载到容器
    environment:
      - MODEL_PATH=/models/llama-13b
      - QUANTIZATION=INT4
      - MAX_SEQ_LEN=4096
      - BATCH_SIZE=4
      - PORT=8080
    ports:
      - "8080:8080"
    restart: unless-stopped

启动服务：

bash
docker-compose up -d

验证部署

服务启动后，可通过HTTP请求验证推理功能：

bash
curl -X POST http://localhost:8080/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "What is AI?", "max_new_tokens": 100}'

预期返回模型生成的文本结果。