ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

Diffusers 量化实战:用 bitsandbytes 8-bit/4-bit 在 16GB 显存内运行 FLUX.1-dev

2026/9/10 11:52:20 拓冰建站 浏览量
Diffusers 量化实战:用 bitsandbytes 8-bit/4-bit 在 16GB 显存内运行 FLUX.1-dev Diffusers 量化实战用 bitsandbytes 8-bit/4-bit 在 16GB 显存内运行 FLUX.1-dev【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本指南系统讲解如何在 Diffusers 中使用 bitsandbytes 库对扩散模型进行 8-bitLLM.int8() 算法与 4-bitQLoRA/NF4 算法量化。通过文中的完整示例你将掌握如何量化 FLUX.1-dev 的 DiT 主干与 T5 文本编码器把模型总显存占用压到 16GB 以下甚至能在免费的 Google Colab 上运行并学会 outlier 阈值调优、模块跳过、嵌套量化、反量化与torch.compile加速等进阶技巧。为什么选择 bitsandbytes在扩散模型的推理与微调中显存往往比算力更先成为瓶颈。以 FLUX.1-dev 为例其 DiT 主干FluxTransformer2DModel与文本编码器T5EncoderModel合计参数规模庞大未经处理时很难装入消费级显卡。bitsandbytes 是目前将模型量化为 8-bit 和 4-bit 的最便捷方案其核心思路是在加载权重时把torch.nn.Linear层替换为 bitsandbytes 提供的低精度线性层从而大幅压缩模型驻留显存8-bitLLM.int8() 算法将超出阈值的异常值outlier以 fp16 单独计算其余非异常值以 int8 计算随后把非异常值的计算结果转回 fp16 并与异常值结果相加最终以 fp16 返回权重。这种做法降低了异常值对模型性能的破坏性影响实测显存占用可减半。4-bitQLoRA 算法在 8-bit 基础上进一步压缩显存占用可降至约四分之一常与 QLoRA 技术结合用于量化 LLM 的微调即仅训练额外参数见后文注意事项。在 Diffusers 中bitsandbytes 同时得到 Diffusers 与 Transformers 两个生态的支持因此你可以分别量化扩散模型的 DiT 主干如FluxTransformer2DModel和基于 Transformers 的文本编码器如T5EncoderModel。环境准备与安装量化功能依赖 bitsandbytes 与 Accelerate 协同工作。安装最新版依赖pip install diffusers transformers accelerate bitsandbytes -U从源码实现看量化器 对环境有明确的硬性要求建议对照检查GPU 是必需的4-bit 量化器要求torch.cuda、torch.xpu或torch.mps至少其一可用否则直接抛出RuntimeError(No GPU found. A GPU is needed for quantization.)8-bit 量化器同样要求 CUDA 或 XPU。bitsandbytes 0.43.38-bit 与 4-bit 量化器在validate_environment中都会校验版本低于该版本会提示pip install -U bitsandbytes。accelerate 0.26.0量化权重加载依赖 Accelerate 的init_empty_weights与设备映射能力。此外BitsAndBytesConfig的post_init还会单独校验 4-bit 量化需要bitsandbytes 0.39.0quantization_config.py。8-bit 量化实战跑通 FLUX.1-dev量化模型的方式很简单把BitsAndBytesConfig传给from_pretrained。它适用于任意模态、任意模型只要该模型支持通过 Accelerate 加载且包含torch.nn.Linear层。由于 Diffusers 和 Transformers 各自导出同名配置类建议用别名区分导入。对 FLUX.1-dev 而言文本编码器 T5 体量巨大需要量化而CLIPTextModel与AutoencoderKL本身尺寸已经很小且AutoencoderKL只有少量torch.nn.Linear层无需量化保持原精度即可。from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig import torch from diffusers import AutoModel from transformers import T5EncoderModel # 8-bit 量化 T5 文本编码器 quant_config TransformersBitsAndBytesConfig(load_in_8bitTrue) text_encoder_2_8bit T5EncoderModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertext_encoder_2, quantization_configquant_config, dtypetorch.float16, ) # 8-bit 量化 DiT 主干 quant_config DiffusersBitsAndBytesConfig(load_in_8bitTrue) transformer_8bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, )需要注意两点关于dtypetorch.nn.Linear之外的模块如torch.nn.LayerNorm默认转换为torch.float16可通过dtype参数改变这些模块的数据类型例如改用torch.float32以获得更高的数值精度代价是更大的显存与更慢的速度transformer_8bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float32, )Ada 及以上架构 GPU 建议在 AdaRTX 40 系列及更新的 GPU 上推荐将dtype与bnb_4bit_compute_dtype统一为torch.bfloat16以发挥 bf16 计算单元的吞吐优势。从源码看8-bit 量化通过 utils.py 中的_replace_with_bnb_linear递归遍历模型把nn.Linear原位替换为bnb.nn.Linear8bitLt并传入has_fp16_weights对应llm_int8_has_fp16_weight与threshold对应llm_int8_threshold等参数替换后的模块还会被强制requires_grad_(False)。若未显式指定dtype量化器会强制回退为torch.float16并打印提示日志。组装管线并生成图像接下来把两个量化模型装回FluxPipeline中推理。设置device_mapauto后Accelerate 会按先 GPU、再 CPU、最后硬盘的顺序自动铺满所有可用设备——硬盘是最慢的选择仅在显存与内存都不足时才会使用。from diffusers import FluxPipeline pipe FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, transformertransformer_8bit, text_encoder_2text_encoder_2_8bit, dtypetorch.float16, device_mapauto, ) pipe_kwargs { prompt: A cat holding a sign that says hello world, height: 1024, width: 1024, guidance_scale: 3.5, num_inference_steps: 50, max_sequence_length: 512, } image pipe(**pipe_kwargs, generatortorch.manual_seed(0)).images[0]参数说明guidance_scale3.5为 FLUX 推荐的 CFG 强度max_sequence_length512控制文本编码的序列长度越大对提示词的语义理解越充分、显存开销也越高torch.manual_seed(0)用于复现结果。当显存足够时也可以直接用.to(cuda)把整条管线搬到 GPU再调用enable_model_cpu_offload开启模型级 CPU 卸载以进一步优化显存占用——CPU offload 会在每个子模块执行前把权重搬到 GPU、执行完再搬回牺牲部分速度换取显存。4-bit 量化实战4-bit 流程与 8-bit 几乎完全一致只是配置改为load_in_4bitTrue显存占用可再降约一半相对 fp16 约为四分之一。from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig import torch from diffusers import AutoModel from transformers import T5EncoderModel # 4-bit 量化 T5 文本编码器 quant_config TransformersBitsAndBytesConfig(load_in_4bitTrue) text_encoder_2_4bit T5EncoderModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertext_encoder_2, quantization_configquant_config, dtypetorch.float16, ) # 4-bit 量化 DiT 主干 quant_config DiffusersBitsAndBytesConfig(load_in_4bitTrue) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, )后续组装管线、生成图像与 8-bit 完全相同将transformer_8bit/text_encoder_2_8bit换成 4-bit 版本即可device_mapauto与enable_model_cpu_offload的使用方式也一致。量化器在底层会把nn.Linear替换为bnb.nn.Linear4bit并传入计算精度bnb_4bit_compute_dtype、compress_statistics嵌套量化开关与quant_type等参数。关于 4-bit 量化的自动精度回退与 8-bit 相同未指定dtype时量化器会强制使用torch.float16这是 bitsandbytes 加载低比特权重的硬性要求。保存与重新加载量化模型模型量化完成后可以用标准的push_to_hub上传到 Hub或save_pretrained本地保存序列化权重。保存时先写入包含量化信息的config.json再写入量化权重8-bit 权重还会附带SCB统计分量4-bit 权重则附带quant_state分量。重新加载时无需再传quantization_config——只要模型目录中的config.json带有quantization_config字段内含quant_method: bitsandbytesfrom_pretrained会自动识别并恢复量化加载。例如从预量化仓库加载 4-bit 模型from diffusers import AutoModel, BitsAndBytesConfig quantization_config BitsAndBytesConfig(load_in_4bitTrue) model_4bit AutoModel.from_pretrained( hf-internal-testing/flux.1-dev-nf4-pkg, subfoldertransformer )从源码看这一加载路径由 bnb_quantizer.py 中的check_if_quantized_param/create_quantized_param实现4-bit 预量化权重通过bnb.nn.Params4bit.from_prequantized结合quant_state恢复8-bit 预量化权重则要求SCB分量齐全且权重为torch.int8。量化器还声明了is_serializable与is_trainable恒为True因为强制要求 bitsandbytes 0.43.3天然支持序列化。分片检查点场景下8-bit 量化器会通过maybe_update_state_dict暂存未配对的 weight/SCB等对应分片到达后再合并。8-bit 特性进阶Outlier 阈值与模块跳过Outlier 阈值llm_int8_thresholdLLM.int8() 的核心是把异常值单独拎出来用 fp16 计算。所谓异常值是指隐藏状态中超过某个阈值的取值。正常权重分布通常在 [-3.5, 3.5]但大模型的分布可能差异极大如 [-60, 6] 或 [6, 60]。int8 对幅值约 5 以内的值表现良好超出后会有明显性能损失。默认阈值 6 是不错的起点但对更不稳定的模型小模型或微调中的模型可能需要调低。建议通过llm_int8_threshold参数实验寻找最优阈值from diffusers import AutoModel, BitsAndBytesConfig quantization_config BitsAndBytesConfig( load_in_8bitTrue, llm_int8_threshold10, ) model_8bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquantization_config, )该参数最终会作为threshold传给bnb.nn.Linear8bitLt见 utils.py源码中其默认值为6.0并校验必须是 floatquantization_config.py。跳过模块转换llm_int8_skip_modules并非所有模块都适合量化成 8-bit强行量化反而可能引入数值不稳定。以 Stable Diffusion 3 的SD3Transformer2DModel为例proj_out最终输出投影层可以跳过转换保持原始精度from diffusers import SD3Transformer2DModel, BitsAndBytesConfig quantization_config BitsAndBytesConfig( load_in_8bitTrue, llm_int8_skip_modules[proj_out], ) model_8bit SD3Transformer2DModel.from_pretrained( stabilityai/stable-diffusion-3-medium-diffusers, subfoldertransformer, quantization_configquantization_config, )跳过模块通过modules_to_not_convert贯穿整个替换过程_process_model_before_weight_loading会把llm_int8_skip_modules与需要保持 fp32 的模块合并utils.py 在递归替换时对命中名单的nn.Linear不做替换。此外该名单还会自动追加被device_map调度到cpu/disk的模块键这些模块会保持 32-bit 而非量化见_process_model_before_weight_loading中对llm_int8_enable_fp32_cpu_offload的校验逻辑。4-bit 特性进阶计算精度、NF4 与嵌套量化计算数据类型bnb_4bit_compute_dtype4-bit 权重存储很省显存但计算时默认以 fp32 进行速度偏慢。将计算类型改为 bf16 可显著提速import torch from diffusers import BitsAndBytesConfig quantization_config BitsAndBytesConfig(load_in_4bitTrue, bnb_4bit_compute_dtypetorch.bfloat16)源码中该参数默认值为torch.float32支持传入字符串或torch.dtype两种形式字符串会被getattr(torch, ...)解析并在post_init校验类型。Normal Float 4NF4量化类型NF4 是 QLoRA 论文提出的一种 4-bit 数据类型针对从正态分布初始化的权重做了适配因此在训练 4-bit 基座模型时应使用 NF4。通过bnb_4bit_quant_typenf4配置from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig from diffusers import AutoModel from transformers import T5EncoderModel quant_config TransformersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_quant_typenf4, ) text_encoder_2_4bit T5EncoderModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertext_encoder_2, quantization_configquant_config, dtypetorch.float16, ) quant_config DiffusersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_quant_typenf4, ) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, )bnb_4bit_quant_type的默认值为fp4可选fp4或nf4。对纯推理而言quant_type对性能影响不大但为了与预量化权重保持一致仍建议同时设置匹配的bnb_4bit_compute_dtype与dtype取值。量化器会通过quantization_method()把 4-bit 配置归类为fp4或nf4见 quantization_config.py。嵌套量化Nested Quantization嵌套量化double quantization对已经量化的权重再做一次量化即对第一次量化产生的量化常数进行二次量化可额外节省约 0.4 bits/参数且几乎不损失性能from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig from diffusers import AutoModel from transformers import T5EncoderModel quant_config TransformersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_use_double_quantTrue, ) text_encoder_2_4bit T5EncoderModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertext_encoder_2, quantization_configquant_config, dtypetorch.float16, ) quant_config DiffusersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_use_double_quantTrue, ) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, )该开关在底层对应bnb.nn.Linear4bit的compress_statistics参数见 utils.py默认关闭。反量化Dequantize量化后的模型可以恢复到原始精度但可能有轻微质量损失且反量化后的模型会重新占用大显存请确保 GPU RAM 足够容纳。调用模型的dequantize()方法即可from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig from diffusers import AutoModel from transformers import T5EncoderModel quant_config TransformersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_use_double_quantTrue, ) text_encoder_2_4bit T5EncoderModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertext_encoder_2, quantization_configquant_config, dtypetorch.float16, ) quant_config DiffusersBitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_use_double_quantTrue, ) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, ) text_encoder_2_4bit.dequantize() transformer_4bit.dequantize()底层实现位于量化器的_dequantize方法若模型处于 CPU例如经历过enable_model_cpu_offload会先临时搬到 GPU通过dequantize_and_replace把Linear4bit/Linear8bitLt还原为普通线性层再按原设备放回。检查内存占用用get_memory_footprint方法查看量化前后模型参数的内存占用print(model.get_memory_footprint())注意该方法只统计模型参数本身的驻留内存不估算推理过程中的峰值显存需求激活值、中间张量与 CUDA 上下文同样占用显存因此实际推理显存会高于该数值。用 torch.compile 加速推理量化省显存、torch.compile提速度二者可以叠加。使用torch.compile需要最新版 bitsandbytes并建议安装 PyTorch nightly# 8-bit torch._dynamo.config.capture_dynamic_output_shape_ops True quant_config DiffusersBitsAndBytesConfig(load_in_8bitTrue) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, ) transformer_4bit.compile(fullgraphTrue)# 4-bit quant_config DiffusersBitsAndBytesConfig(load_in_4bitTrue) transformer_4bit AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, ) transformer_4bit.compile(fullgraphTrue)在 8-bit 场景下需要先开启torch._dynamo.config.capture_dynamic_output_shape_ops True才能成功捕获动态输出形状算子fullgraphTrue则要求整张计算图被完整编译若中间存在不受支持的算子会直接报错。从量化器实现看8-bit 量化器显式声明了is_compileable True即官方支持对量化后模型做图编译。参考基准数据RTX 4090开启编译后4-bit FLUX 生成时间从无编译的 32.570 秒降至 25.809 秒。关于训练与微调的注意点重要警告8-bit 与 4-bit 权重的训练仅支持训练额外参数即 LoRA/Adapter 等旁路参数基座权重本身不可训练。若要在量化模型之上微调 LoRA可参考 HiDream 训练示例中的量化章节其中展示了在量化 Transformer 之上附加可训练 LoRA 层的完整流程是 QLoRA 思路在扩散模型上的落地范本。小结8-bitLLM.int8()适合追求更低精度损失的快速部署显存减半4-bitQLoRA/NF4适合极致省显存或微调场景显存约为四分之一。BitsAndBytesConfig同时存在于 Diffusers 与 Transformers分别用于量化 DiT 主干与文本编码器。高频调优参数集中在 quantization_config.py8-bit 关注llm_int8_threshold与llm_int8_skip_modules4-bit 关注bnb_4bit_compute_dtype、bnb_4bit_quant_type与bnb_4bit_use_double_quant。量化加载、反量化、序列化与设备映射的底层逻辑可在 bnb_quantizer.py 与 utils.py 中追溯相关行为由 tests/quantization/bnb/ 下的 4-bit 与混合 int8 测试覆盖验证。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考