neural-compressor

环境要求

• Python 环境（用于 pip 安装 neural-compressor）
• PyTorch 后端：CPU/GPU 需安装 intel_extension_for_pytorch，其他平台安装常规 torch
• 运行 FP8 示例需要 Intel Gaudi2 AI 加速器
• 使用 Gaudi 时需使用带 Gaudi 软件栈的 Docker 镜像，并保证与 Neural Compressor 版本匹配

安装与启动步骤

1. 1准备 PyTorch 框架

   按硬件环境选择安装：CPU 版或 Intel GPU 版 intel_extension_for_pytorch，其他平台直接安装 torch。

2. 2安装 Neural Compressor

   从 PyPI 安装该库，建议在已装好 PyTorch 的虚拟环境中执行，避免依赖冲突。

   pip install neural-compressor

   复制

3. 3启动 Gaudi 容器

   跑 FP8 示例推荐使用 Gaudi 软件栈镜像，以交互式 shell 启动容器并挂载全部 Gaudi 设备。

   docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.24.0/ubuntu24.04/habanalabs/pytorch-installer-2.10.0:latest

   复制

4. 4设置懒执行模式

   Habana 1.21.0 起默认为 0，但多数低精度函数不支持，建议设为 1 以保证兼容。

   export PT_HPU_LAZY_MODE=1

   复制

5. 5运行 FP8 量化示例

   用 resnet18 演示 FP8 E4M3 量化流程：先 prepare 再 convert，需在 Gaudi2 环境执行。

   python fp8_example.py

   复制

关键配置

如何确认成功

容器内可成功导入 neural_compressor.torch.quantization，且示例中 prepare/convert 执行无报错。

常见问题

注意事项

• Neural Compressor 与 Gaudi 软件栈存在版本映射，务必使用匹配组合。
• Gaudi 上跑低精度功能建议保持 PT_HPU_LAZY_MODE=1。
• 非 Gaudi 用户无需启动 Habana 容器，直接安装对应 PyTorch 版本即可。
• 安装 neural-compressor 前请先完成 PyTorch 后端依赖安装。