vLLM 中的 TorchAO 量化从模型量化保存到 vLLM 推理加载的完整实践【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllmTorchAO 是 PyTorch 官方的架构优化库提供面向推理与训练的高性能 dtypes、优化技术与内核并能与 torch.compile、FSDP 等原生 PyTorch 能力组合使用。vLLM 通过内置的torchao量化方法接入该库既支持直接加载已经用 TorchAO 量化并保存的 HuggingFace 检查点也支持在加载 bf16 原始权重时进行在线量化online quantization。读完本文你将掌握用 TorchAO 量化自有模型并推送到 Hub 的方法、vLLM 加载 TorchAO 检查点的配置方式以及 vLLM 源码中权重在线量化、硬件打包转换与安全张量safetensors加载策略的底层实现细节。TorchAO 是什么vLLM 如何接入TorchAOtorchao定位是 PyTorch 的架构优化库其核心能力包括提供一系列量化/低精度 dtypes 与张量子类如 Int8/Int4 权重、FP8 权重与激活等提供量化与反量化的算子内核并针对不同硬件如 Hopper 架构上的 fbgemm_gpu_genai preshuffled 内核做优化设计上强调与 PyTorch 原生特性torch.compile、FSDP 等的可组合性。vLLM 在量化配置注册表中将torchao映射到TorchAOConfig类见 vllm/model_executor/layers/quantization/init.py因此只要给 vLLM 传入quantizationtorchao引擎就会用 TorchAO 的张量子类与内核来承载线性层的权重与计算。安装 TorchAO官方推荐使用 TorchAO 的 nightly 版本并选择与系统匹配的 CUDA 版本cu126、cu128 等# 安装最新的 TorchAO nightly 构建 # 注意选择与你系统匹配的 CUDA 版本cu126, cu128, etc. pip install \ --pre torchao10.0.0 \ --index-url https://download.pytorch.org/whl/nightly/cu126需要注意的是TorchAO nightly 目前仅与对应的 PyTorch nightly 组合使用效果最佳。从 vLLM 测试代码 tests/quantization/test_torchao.py 可以看到若干依赖较新 TorchAO 版本0.14.0.dev 及以上的用例因 nightly 兼容性问题被跳过这正是这一限制的佐证。若使用已发布的稳定版本vLLM 侧对版本的要求分散在不同功能点加载 TorchAO 序列化的 safetensors 检查点需要torchao 0.15.0见下文 safetensors 策略配置反序列化torchao.core.config.config_from_dictvLLM 要求torchao 0.10.0否则会抛出带安装提示的 ImportErrorvllm/model_executor/layers/quantization/torchao.py。用 TorchAO 量化 HuggingFace 模型你可以用 TorchAO 量化自己的 HuggingFace 模型transformers、diffusers 均提供TorchAoConfig接口并把量化后的检查点保存到 HuggingFace Hub。例如将 Llama-3-8B 做 Int8 权重量化import torch from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer from torchao.quantization import Int8WeightOnlyConfig model_name meta-llama/Meta-Llama-3-8B quantization_config TorchAoConfig(Int8WeightOnlyConfig()) quantized_model AutoModelForCausalLM.from_pretrained( model_name, dtypeauto, device_mapauto, quantization_configquantization_config ) tokenizer AutoTokenizer.from_pretrained(model_name) input_text What are we having for dinner? input_ids tokenizer(input_text, return_tensorspt).to(cuda) hub_repo # YOUR HUB REPO ID tokenizer.push_to_hub(hub_repo) quantized_model.push_to_hub(hub_repo, safe_serializationFalse)要点说明Int8WeightOnlyConfig()只做权重量化权重 int8、激活保持原精度属于通用性最好、对硬件要求最低的配置TorchAO 还提供 FP8 权重、FP8 动态激活 FP8 权重、Int4 权重等更多Config类型可替换TorchAoConfig(...)中的参数推送到 Hub 时使用safe_serializationFalse本例采用 .pt 格式若要走 safetensors 分片格式保存则需要在加载侧配合 vLLM 的torchao加载策略见后文除了手写脚本TorchAO 官方还提供了一个带简单 UI 的量化 Space适合不想写代码的用户批量量化模型。量化并推送后的检查点在config.json中会带有quant_method: torchao、quant_type: {default: {...}}等字段——这正是 vLLM 识别 TorchAO 检查点的依据。vLLM 如何识别与解析 TorchAO 检查点在 vllm/model_executor/layers/quantization/torchao.py 中TorchAOConfig.from_config定义了 vLLM 对检查点config.json的解析规则读取quant_method字段若其包含字符串torchao则标记is_checkpoint_torchao_serializedTrue表示检查点本身已是 TorchAO 量化格式读取quant_type字段必须存在且只含一个键default其值为 TorchAO 配置的 dict通过torchao.core.config.config_from_dict还原成 TorchAO 的 config 对象汇总跳过列表config.json中的modules_to_not_convert以及quant_type的_data.module_fqn_to_config中值为None的模块都会进入skip_modules。基础能力方面同一个文件get_name()返回torchao即命令行/LLM参数中使用的量化方法名get_supported_act_dtypes()返回[float32, float16, bfloat16]即激活侧 dtype 的适用范围get_min_capability()返回75表示最低支持 Turingsm75及更高计算能力。对于 FP8 激活类配置名称中含Float8与Activation的配置如Float8DynamicActivationFloat8WeightConfigvLLM 会在量化前做前置检查_check_torchao_fp8_activation_capabilitytorchao.py要求 GPU 计算能力 8.9NVIDIA Ada/Hopper或 AMD MI300否则直接抛出带明确指引的错误建议在老 GPU 上改用Int8WeightOnlyConfig或Int4WeightOnlyConfig。两种加载模式预量化检查点 vs 在线量化模式一加载预量化检查点直接传入已量化的模型即可from vllm import LLM llm LLM( modelyour-org/llama3-8b-int8wo, # 你用 TorchAO 量化并推送的模型 quantizationtorchao, )vLLM 测试 tests/quantization/test_torchao.py 中的test_pre_quantized_model即加载torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.15.0这类预量化检查点并验证前向输出有限。模式二在线量化Online Quantization如果手上只有 bf16/fp16 原始权重可以借助hf_overrides把 TorchAO 配置注入 vLLM让 vLLM 在加载完成后现场量化。TorchAO 配置有两条注入通道对应TorchAOConfig.from_config_dict_json与TorchAOConfig.from_config_file两个入口quantization_config_dict_json把 TorchAO config 对象json.dumps(config_to_dict(config))后的字符串直接写进hf_overridesquantization_config_file把该 JSON 写到本地文件传文件路径。测试用例test_online_quant_config_dict_json/test_online_quant_config_file给出的完整用法加载facebook/opt-125m并在线量化为 FP8import json from torchao.core.config import config_to_dict from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow from vllm import LLM torchao_quant_config Float8DynamicActivationFloat8WeightConfig(granularityPerRow()) hf_overrides { quantization_config_dict_json: json.dumps(config_to_dict(torchao_quant_config)) } llm LLM( modelfacebook/opt-125m, dtypebfloat16, quantizationtorchao, hf_overrideshf_overrides, enforce_eagerTrue, )此外还支持按模块细粒度配置当 TorchAO 配置是ModuleFqnToConfig时get_quant_method会按模块全限定名FQN精确匹配或匹配re:前缀的正则模式未命中则回退_default某模块没有配置时返回UnquantizedLinearMethod即该层保持不量化。配合skip_modules来源为modules_to_not_convert或 FQN 配置中值为None的项可以实现部分量化——测试test_opt_125m_int8wo_model_loading_with_params加载的jerryzh168/opt-125m-int8wo-partial-quant模型即为此类场景。在线量化的源码级实现TorchAOLinearMethodtorchao.py是 vLLM 中 TorchAO 线性层的计算方法其工作流为create_weights按 vLLM 的张量并行切分形状创建原始 dtype 的权重Parameter标记input_dim1, output_dim0供weight_loader切分写入。若检查点已是 TorchAO 序列化格式则直接在 meta 设备上构造量化后的权重占位避免真实显存分配process_weights_after_loading权重加载完成后触发量化。核心函数是torchao_quantize_param_data——它在 meta 设备上构造一个包裹nn.Linear的nn.Sequential注释说明quantize_是原地操作且某些配置需要模块替换只有非顶层模块才支持模块替换因此必须包一层把真实权重挂上去后调用 TorchAO 的quantize_完成量化返回量化后的张量子类硬件适配若torchao 0.15.0随后调用convert_to_packed_tensor_based_on_current_hardware把量化张量转换为适合当前硬件的打包格式。例如在装有fbgemm_gpu_genai且 GPU 为 sm90H100时Int4 权重会从普通的Int4Tensor转换为Int4PreshuffledTensor以便走 preshuffled GEMM 内核——测试test_opt_125m_int4wo_model_running_preshuffled_kernel专门验证了该转换行为apply前向就是朴素的F.linear(x, layer.weight, bias)反量化与低精度 GEMM 全部由 TorchAO 张量子类在内部完成。权重属性如weight_loader绑定方法、input_dim/output_dim在张量替换前后会通过_get_weight_attrs/_restore_weight_attrs记录并恢复保证并行加载钩子不被量化替换破坏。测试 tests/quantization/test_torchao.py 中的test_reload_weights还展示了进阶用法以load_formatdummy启动权重为随机初始化运行collective_rpc(reload_weights)后按 TorchAO 配置原地重载真实权重并在线量化——这对需要频繁换权重的 RLHF/RL 场景有价值vLLM 在 vllm/model_executor/model_loader/reload/torchao_decorator.py 中为其提供了适配。用 safetensors 分片保存的 TorchAO 检查点TorchAO 张量子类在 safetensors 中会被展平为普通张量并按元数据记录结构无法直接按模型参数加载因此 vLLM 在SafetensorsLoadStrategy中增加了专门的torchao策略vllm/config/load.pySafetensorsLoadStrategy: TypeAlias Literal[lazy, eager, prefetch, torchao]其实现位于 vllm/model_executor/model_loader/weight_utils.py该策略要求torchao 0.15.0否则抛出错误逐分片safe_open读取张量与 metadata调用torchao.prototype.safetensors.safetensors_support.unflatten_tensor_state_dict把展平的张量重组回 TorchAO 张量子类由于分片检查点中同一张量子类的数据可能分布在多个文件中该函数会返回本步无法完成的剩余数据leftover_state_dict在后续分片迭代中补齐——这是对多分点场景的关键处理。对应地DefaultModelLoadervllm/model_executor/model_loader/default_loader.py会根据load_config.safetensors_load_strategy选择该迭代器。也就是说如果你用push_to_hub以 safetensors 分片形式保存了 TorchAO 量化模型加载时应确认 vLLM 选择了torchao加载策略且环境中torchao 0.15.0。适用前提与限制小结场景前提/限制依据通用 TorchAO 量化方法GPU 计算能力 sm75激活 dtype 支持 fp32/fp16/bf16TorchAOConfig.get_min_capability、get_supported_act_dtypesFP8 激活类配置如Float8DynamicActivationFloat8WeightConfigGPU 计算能力 8.9Ada/Hopper或 AMD MI300_check_torchao_fp8_activation_capability从config.json反序列化 TorchAO 配置torchao 0.10.0TorchAOConfig.from_config的 ImportError 提示safetensors 加载 TorchAO 检查点safetensors_load_strategytorchaotorchao 0.15.0weight_utils.pyInt4 preshuffled 内核Hoppersm90且安装fbgemm_gpu_genai否则自动回退普通Int4Tensorconvert_to_packed_tensor_based_on_current_hardware、test_opt_125m_int4wo_model_running_preshuffled_kernelnightly 版本TorchAO nightly 与 PyTorch nightly 配套测试文件中多处pytest.mark.skip原因说明从源码结构看vLLM 的 TorchAO 集成是一个薄适配层配置解析、模块跳过、在线量化、硬件打包、safetensors 重组都复用了 TorchAO 自身 APIvLLM 主要负责把量化生命周期嵌入到模型加载流程create_weights→ 权重加载 →process_weights_after_loading之中。这也意味着你可以直接使用 TorchAO 社区提供的任意AOBaseConfig子类只要它满足硬件前提vLLM 侧不需要额外支持代码。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考