diffusers 如何用 device_map 在多个 GPU 间分发大型模型组件【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers当单个 GPU 显存装不下扩散模型例如参数超过 100 亿的 FLUX.1-dev transformer时diffusers 提供的device_map参数可以把 pipeline 的不同组件、或单个模型的不同层自动分配到多张 GPU 上让大模型在消费级显卡上也能够完成推理。这篇文章基于仓库中的 Distributed inference 指南 和 Reduce memory usage 指南给出用device_map在多 GPU 间分发大型模型组件的完整操作路径加载、验证分布结果、按阶段加载/卸载组件、以及完成一张图的生成。准备条件device_map相关能力由 Accelerate 库支撑需要先安装它见 Reduce memory usage 指南pip install -U accelerate操作对象是两块或多块 GPU。文档中所有多卡示例都以两块 GPU 为前提max_memory{0: 16GB, 1: 16GB}表示 GPU 0 和 GPU 1。需要注意一个官方警告device placement 是实验性功能API 可能会变化pipeline 级别目前只支持balanced策略未来计划支持更多映射策略。Pipeline 级分发device_mapbalanceddevice_map参数控制 pipeline 内各组件在设备间的分布。在 pipeline 上支持两个取值见 加载指南参数说明cuda把整个 pipeline 放到一个受支持的加速器设备上balanced把 pipeline 组件均匀分布到所有 GPU用device_mapbalanced加载 pipeline并用max_memory限制每张 GPU 上允许分配的最大显存import torch from diffusers import StableDiffusionXLPipeline pipeline StableDiffusionXLPipeline.from_pretrained( stabilityai/stable-diffusion-xl-base-1.0, dtypetorch.float16, device_mapbalanced )max_memory传一个字典把每个设备的最大显存用量限定住。不在max_memory中列出的设备会被忽略组件不会被分发到它上面如果不传Diffusers 默认使用所有设备的最大可用显存max_memory {0: 1GB, 1: 1GB} pipeline StableDiffusionXLPipeline.from_pretrained( stabilityai/stable-diffusion-xl-base-1.0, dtypetorch.float16, device_mapbalanced, max_memorymax_memory )用 hf_device_map 验证分布结果加载完成后通过hf_device_map属性查看组件实际落在哪张卡上。下面是 Reduce memory usage 指南 中给出的文档示例输出SDXL pipeline 的balanced分布print(pipeline.hf_device_map) # {unet: 1, vae: 1, safety_checker: 0, text_encoder: 0}这段输出仅表示该示例中各组件的分发结果具体数值取决于你的模型和显存配置不要把它当作固定预期。对单个模型例如 transformer也可以直接调用hf_device_map查看层级的分布。模型级分发device_mapauto 与自定义字典对特别大的单个模型文档以 12.5B 参数的 Flux diffusion transformer 为例在模型级别使用device_mapauto该策略会优先把模型放到最快的设备上放不下时再放到 CPU 或磁盘等更慢的设备上。代价是部分参数存放在慢速设备上时推理延迟会变高import torch from diffusers import AutoModel transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapauto, dtypetorch.bfloat16 )也可以设计自己的device_map。它是一个字典键是模型的具体模块名或层名值是设备标识符GPU 用整数CPU 用cpu磁盘用disk。设计步骤是先调用hf_device_map查看auto策略下的层级分布再基于它调整自己的映射。例如把single_transformer_blocks.10到single_transformer_blocks.20放到第二张 GPU1import torch from diffusers import AutoModel device_map { pos_embed: 0, time_text_embed: 0, context_embedder: 0, x_embedder: 0, transformer_blocks: 0, single_transformer_blocks.0: 0, # ... 省略至 .9 均为 0 ... single_transformer_blocks.10: 1, single_transformer_blocks.11: 1, # ... .11 - .20 均为 1 ... single_transformer_blocks.21: cpu, # 其余 block 放 cpu norm_out: cpu, proj_out: cpu } transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapdevice_map, dtypetorch.bfloat16 )字典中未列出的模块不要自行猜测归属先打印hf_device_map拿到完整映射再在其基础上修改。完整操作路径FLUX.1-dev 分阶段加载与卸载Distributed inference 指南 给出了一个在两块 GPU 上运行 FLUX.1-dev 的完整示例核心思路是按阶段只加载当前需要的组件并把用完的组件从显存中移除。示例假设两块 GPU 可用。第 1 步用 balanced 策略加载文本编码器跳过其他组件。设置device_mapbalanced把文本编码器均匀分布到所有 GPU用max_memory限定每张卡的用量transformerNone和vaeNone表示这两个组件不加载避免占用显存from diffusers import FluxPipeline import torch prompt cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain pipeline FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, transformerNone, vaeNone, device_mapbalanced, max_memory{0: 16GB, 1: 16GB}, dtypetorch.bfloat16 ) with torch.no_grad(): print(Encoding prompts.) prompt_embeds, pooled_prompt_embeds, text_ids pipeline.encode_prompt( promptprompt, prompt_2None, max_sequence_length512 )第 2 步文本编码完成后把编码器组件从 GPU 上移除为 diffusion transformer 腾出空间import gc def flush(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() del pipeline.text_encoder del pipeline.text_encoder_2 del pipeline.tokenizer del pipeline.tokenizer_2 del pipeline flush()第 3 步用 auto 策略把 transformer 自动分布到两块 GPUfrom diffusers import AutoModel import torch transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapauto, dtypetorch.bfloat16 )此时可以用transformer.hf_device_map查看 transformer 本身在设备间的分布。第 4 步把 transformer 装回 pipeline其他组件全部传 None生成 latents 而不是直接解码图像pipeline FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, text_encoderNone, text_encoder_2None, tokenizerNone, tokenizer_2None, vaeNone, transformertransformer, dtypetorch.bfloat16 ) print(Running denoising.) height, width 768, 1360 latents pipeline( prompt_embedsprompt_embeds, pooled_prompt_embedspooled_prompt_embeds, num_inference_steps50, guidance_scale3.5, heightheight, widthwidth, output_typelatent, ).images第 5 步移除 pipeline 和 transformer单独加载 VAE 解码 latents。VAE 通常足够小可以放在单卡上。下面代码中的ckpt_id是文档示例的变量替换为与前面相同的模型 idblack-forest-labs/FLUX.1-devimport torch from diffusers import AutoencoderKL from diffusers.image_processor import VaeImageProcessor ckpt_id black-forest-labs/FLUX.1-dev vae AutoencoderKL.from_pretrained(ckpt_id, subfoldervae, dtypetorch.bfloat16).to(cuda) # or mps, xpu, cpu vae_scale_factor 2 ** (len(vae.config.block_out_channels) - 1) image_processor VaeImageProcessor(vae_scale_factorvae_scale_factor) with torch.no_grad(): print(Running decoding.) latents FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor) latents (latents / vae.config.scaling_factor) vae.config.shift_factor image vae.decode(latents, return_dictFalse)[0] image image_processor.postprocess(image, output_typepil) image[0].save(split_transformer.png)生成出split_transformer.png即说明这条分阶段加载、卸载、再加载的链路跑通。限制与后续切换切换回单卡或 offload 之前必须先reset_device_map对已经做过 device mapping 的 pipeline如果想再用.to()、enable_sequential_cpu_offload或enable_model_cpu_offload需要先调用pipeline.reset_device_map()模型放不下所有 GPU 时文档说明此时需要退回到单卡方案配合enable_model_cpu_offload仅限单卡或enable_sequential_cpu_offload仅限单卡且非常慢详见 Reduce memory usage 指南 的 Offloading 章节。何时该用 device_mapDistributed inference 指南 的策略对比表给出的判断是——当模型装不下、且 GPU 间互连带宽较慢时device_map按组件分发降低权重显存单 prompt 延迟略差是合适选择如果互连带宽快如 NVLink单个组件权重放不下时 tensor parallelism 更优如果装不下的是激活值长序列、高分辨率、视频而不是权重应使用 context parallelism。实验性pipeline 级 device placement 目前仅支持balanced策略API 可能变化生产使用前注意版本行为。【免费下载链接】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），仅供参考 diffusers 如何用 device_map 在多个 GPU 间分发大型模型组件
diffusers 如何用 device_map 在多个 GPU 间分发大型模型组件【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers当单个 GPU 显存装不下扩散模型例如参数超过 100 亿的 FLUX.1-dev transformer时diffusers 提供的device_map参数可以把 pipeline 的不同组件、或单个模型的不同层自动分配到多张 GPU 上让大模型在消费级显卡上也能够完成推理。这篇文章基于仓库中的 Distributed inference 指南 和 Reduce memory usage 指南给出用device_map在多 GPU 间分发大型模型组件的完整操作路径加载、验证分布结果、按阶段加载/卸载组件、以及完成一张图的生成。准备条件device_map相关能力由 Accelerate 库支撑需要先安装它见 Reduce memory usage 指南pip install -U accelerate操作对象是两块或多块 GPU。文档中所有多卡示例都以两块 GPU 为前提max_memory{0: 16GB, 1: 16GB}表示 GPU 0 和 GPU 1。需要注意一个官方警告device placement 是实验性功能API 可能会变化pipeline 级别目前只支持balanced策略未来计划支持更多映射策略。Pipeline 级分发device_mapbalanceddevice_map参数控制 pipeline 内各组件在设备间的分布。在 pipeline 上支持两个取值见 加载指南参数说明cuda把整个 pipeline 放到一个受支持的加速器设备上balanced把 pipeline 组件均匀分布到所有 GPU用device_mapbalanced加载 pipeline并用max_memory限制每张 GPU 上允许分配的最大显存import torch from diffusers import StableDiffusionXLPipeline pipeline StableDiffusionXLPipeline.from_pretrained( stabilityai/stable-diffusion-xl-base-1.0, dtypetorch.float16, device_mapbalanced )max_memory传一个字典把每个设备的最大显存用量限定住。不在max_memory中列出的设备会被忽略组件不会被分发到它上面如果不传Diffusers 默认使用所有设备的最大可用显存max_memory {0: 1GB, 1: 1GB} pipeline StableDiffusionXLPipeline.from_pretrained( stabilityai/stable-diffusion-xl-base-1.0, dtypetorch.float16, device_mapbalanced, max_memorymax_memory )用 hf_device_map 验证分布结果加载完成后通过hf_device_map属性查看组件实际落在哪张卡上。下面是 Reduce memory usage 指南 中给出的文档示例输出SDXL pipeline 的balanced分布print(pipeline.hf_device_map) # {unet: 1, vae: 1, safety_checker: 0, text_encoder: 0}这段输出仅表示该示例中各组件的分发结果具体数值取决于你的模型和显存配置不要把它当作固定预期。对单个模型例如 transformer也可以直接调用hf_device_map查看层级的分布。模型级分发device_mapauto 与自定义字典对特别大的单个模型文档以 12.5B 参数的 Flux diffusion transformer 为例在模型级别使用device_mapauto该策略会优先把模型放到最快的设备上放不下时再放到 CPU 或磁盘等更慢的设备上。代价是部分参数存放在慢速设备上时推理延迟会变高import torch from diffusers import AutoModel transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapauto, dtypetorch.bfloat16 )也可以设计自己的device_map。它是一个字典键是模型的具体模块名或层名值是设备标识符GPU 用整数CPU 用cpu磁盘用disk。设计步骤是先调用hf_device_map查看auto策略下的层级分布再基于它调整自己的映射。例如把single_transformer_blocks.10到single_transformer_blocks.20放到第二张 GPU1import torch from diffusers import AutoModel device_map { pos_embed: 0, time_text_embed: 0, context_embedder: 0, x_embedder: 0, transformer_blocks: 0, single_transformer_blocks.0: 0, # ... 省略至 .9 均为 0 ... single_transformer_blocks.10: 1, single_transformer_blocks.11: 1, # ... .11 - .20 均为 1 ... single_transformer_blocks.21: cpu, # 其余 block 放 cpu norm_out: cpu, proj_out: cpu } transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapdevice_map, dtypetorch.bfloat16 )字典中未列出的模块不要自行猜测归属先打印hf_device_map拿到完整映射再在其基础上修改。完整操作路径FLUX.1-dev 分阶段加载与卸载Distributed inference 指南 给出了一个在两块 GPU 上运行 FLUX.1-dev 的完整示例核心思路是按阶段只加载当前需要的组件并把用完的组件从显存中移除。示例假设两块 GPU 可用。第 1 步用 balanced 策略加载文本编码器跳过其他组件。设置device_mapbalanced把文本编码器均匀分布到所有 GPU用max_memory限定每张卡的用量transformerNone和vaeNone表示这两个组件不加载避免占用显存from diffusers import FluxPipeline import torch prompt cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain pipeline FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, transformerNone, vaeNone, device_mapbalanced, max_memory{0: 16GB, 1: 16GB}, dtypetorch.bfloat16 ) with torch.no_grad(): print(Encoding prompts.) prompt_embeds, pooled_prompt_embeds, text_ids pipeline.encode_prompt( promptprompt, prompt_2None, max_sequence_length512 )第 2 步文本编码完成后把编码器组件从 GPU 上移除为 diffusion transformer 腾出空间import gc def flush(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() del pipeline.text_encoder del pipeline.text_encoder_2 del pipeline.tokenizer del pipeline.tokenizer_2 del pipeline flush()第 3 步用 auto 策略把 transformer 自动分布到两块 GPUfrom diffusers import AutoModel import torch transformer AutoModel.from_pretrained( black-forest-labs/FLUX.1-dev, subfoldertransformer, device_mapauto, dtypetorch.bfloat16 )此时可以用transformer.hf_device_map查看 transformer 本身在设备间的分布。第 4 步把 transformer 装回 pipeline其他组件全部传 None生成 latents 而不是直接解码图像pipeline FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, text_encoderNone, text_encoder_2None, tokenizerNone, tokenizer_2None, vaeNone, transformertransformer, dtypetorch.bfloat16 ) print(Running denoising.) height, width 768, 1360 latents pipeline( prompt_embedsprompt_embeds, pooled_prompt_embedspooled_prompt_embeds, num_inference_steps50, guidance_scale3.5, heightheight, widthwidth, output_typelatent, ).images第 5 步移除 pipeline 和 transformer单独加载 VAE 解码 latents。VAE 通常足够小可以放在单卡上。下面代码中的ckpt_id是文档示例的变量替换为与前面相同的模型 idblack-forest-labs/FLUX.1-devimport torch from diffusers import AutoencoderKL from diffusers.image_processor import VaeImageProcessor ckpt_id black-forest-labs/FLUX.1-dev vae AutoencoderKL.from_pretrained(ckpt_id, subfoldervae, dtypetorch.bfloat16).to(cuda) # or mps, xpu, cpu vae_scale_factor 2 ** (len(vae.config.block_out_channels) - 1) image_processor VaeImageProcessor(vae_scale_factorvae_scale_factor) with torch.no_grad(): print(Running decoding.) latents FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor) latents (latents / vae.config.scaling_factor) vae.config.shift_factor image vae.decode(latents, return_dictFalse)[0] image image_processor.postprocess(image, output_typepil) image[0].save(split_transformer.png)生成出split_transformer.png即说明这条分阶段加载、卸载、再加载的链路跑通。限制与后续切换切换回单卡或 offload 之前必须先reset_device_map对已经做过 device mapping 的 pipeline如果想再用.to()、enable_sequential_cpu_offload或enable_model_cpu_offload需要先调用pipeline.reset_device_map()模型放不下所有 GPU 时文档说明此时需要退回到单卡方案配合enable_model_cpu_offload仅限单卡或enable_sequential_cpu_offload仅限单卡且非常慢详见 Reduce memory usage 指南 的 Offloading 章节。何时该用 device_mapDistributed inference 指南 的策略对比表给出的判断是——当模型装不下、且 GPU 间互连带宽较慢时device_map按组件分发降低权重显存单 prompt 延迟略差是合适选择如果互连带宽快如 NVLink单个组件权重放不下时 tensor parallelism 更优如果装不下的是激活值长序列、高分辨率、视频而不是权重应使用 context parallelism。实验性pipeline 级 device placement 目前仅支持balanced策略API 可能变化生产使用前注意版本行为。【免费下载链接】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），仅供参考