摘要

图像生成属于计算机视觉前沿方向，以 Stable Diffusion 为代表的扩散模型广泛应用于文创设计、素材生成、AI 绘图等业务。通用预训练生成模型无法适配行业专属风格、定制物体生成需求，需要通过微调来注入领域知识。昇腾 NPU 平台依托 PyTorch‑NPU 适配层、CANN 算子库、HCCL 集合通信，支持全参数微调、LoRA 轻量微调等图像生成训练方案。图像生成微调显存开销巨大，昇腾通过梯度检查点、混合精度、模型分片、内存优化策略，降低显存占用，完成文生图扩散模型的领域适配微调。本文介绍图像生成微调业务痛点、昇腾平台微调整体流程、LoRA 微调代码示例、环境配置、性能调优与常见问题，帮助开发者在昇腾设备上快速完成图像生成 CV 模型微调落地。

1 背景介绍

扩散类图像生成模型包含文本编码器、UNet 扩散网络、VAE 解码器三大部分，参数量大，训练迭代开销高。直接使用开源预训练模型，生成效果受训练数据集限制，面对企业定制画风、特定产品、IP 形象生成，必须做模型微调。

完整全参数微调需要更新 UNet 全部权重，显存消耗极高；LoRA 微调仅训练少量低秩矩阵，冻结主干网络权重，训练速度更快、显存占用更低，是图像生成场景首选微调方案。

将图像生成微调业务迁移至昇腾平台会遇到几类典型问题：

第一，扩散模型循环迭代步数多，大量小算子，部分算子性能适配不足；

第二，微调训练显存压力巨大，容易出现 OOM 显存溢出；

第三，混合精度、梯度检查点、分布式多卡训练需要适配 NPU 设备；

第四，微调后权重导出、推理部署链路需要完整打通。

昇腾 CANN+PyTorch‑NPU 完整支持 Hugging Face Diffusers 训练链路，无需大规模修改业务代码，就可以完成图像生成模型 LoRA 微调，训练完成权重可以直接用于推理或者导出 om 离线模型部署到 MindIE 推理引擎。

2 昇腾图像生成微调整体业务流程

昇腾平台图像生成微调完整链路分为数据集准备、环境适配、模型微调、权重导出、效果验证五大环节。

1. 数据集准备：整理图文配对数据集，图片 + prompt 文本描述，按照 Diffusers 标准数据集格式组织；

2. 昇腾环境初始化：加载 CANN 环境，PyTorch‑NPU 后端将训练算子调度到 NPU AI Core 执行，HCCL 负责多卡分布式集合通信；

3. LoRA 微调训练：冻结预训练 SD 主干权重，在 UNet 层插入 LoRA 低秩矩阵；开启梯度检查点、FP16 混合精度降低显存；支持单机单卡、单机多卡训练；

4. 权重保存导出：保存 LoRA 权重，支持 safetensors 安全格式；可将 LoRA 权重合并进基础模型；

5. 微调后推理验证：Diffusers pipeline 加载基础模型 + LoRA 权重，NPU 上完成文生图生成，校验生成图像风格、内容是否符合微调预期。

技术选型建议：业务优先选择 LoRA 轻量微调；只有大规模行业数据集场景，才考虑完整全参数微调。

3 关键调优手段，适配昇腾 NPU 图像生成微调

3.1 显存优化策略

• Gradient Checkpoint 梯度检查点：牺牲少量计算，大幅降低显存，是 SD 微调必开配置；

• FP16/BF16 混合精度训练：PyTorch‑NPU 支持 AMP 自动混合精度，减少张量显存占用；

• LoRA 低秩矩阵：仅训练少量参数，相比全量微调显存下降 70% 以上；

• 梯度累积：小 batch 模拟大 batch 效果，降低单步 batch 大小，规避 OOM。

3.2 分布式训练

多卡微调依赖 HCCL 集合通信，DDP 分布式数据并行，多 NPU 分担训练 batch，缩短训练时长。

3.3 算子适配

PyTorch‑NPU 已经适配 Diffusers 绝大多数算子；少量未适配算子会 fallback 到 CPU 执行，造成性能下降，可通过日志定位，升级 CANN 与 PyTorch‑NPU 版本解决。

4 代码实践：昇腾 NPU 上 SD LoRA 图像生成微调

4.1 环境准备脚本 env.sh

# 加载CANN驱动环境
source /usr/local/Ascend/ascend‑toolkit/set_env.sh

# 安装依赖，PyTorch‑NPU，diffusers，accelerate
pip install torch‑npu
pip install diffusers==0.18.1 accelerate transformers datasets safetensors pillow

# NPU混合精度、内存优化环境变量
export ASCEND_GLOBAL_LOG_LEVEL=3
export TORCH_NPU_ALLOC_CONF=expandable_segments:True
# HCCL多卡通信配置，单卡可忽略
export HCCL_CONNECT_TIMEOUT=180

4.2 LoRA 微调核心 Python 代码 train_lora_sd_npu.py

import torch
import os
from datasets import load_dataset
from diffusers import StableDiffusionPipeline
from diffusers.models.lora import LoRALinearLayer
from torch_npu.utils import npu_set_amp

def main():
    device = torch.device("npu:0")
    pretrained_model_name_or_path = "./stable‑diffusion‑v1‑5"
    dataset_path = "./train_dataset"
    output_dir = "./sd_lora_output"

    # 1.加载预训练SD pipeline迁移到NPU
    pipe = StableDiffusionPipeline.from_pretrained(
        pretrained_model_name_or_path, torch_dtype=torch.float16
    ).to(device)

    # 冻结主干UNet、VAE、文本编码器，只训练LoRA权重
    for param in pipe.unet.parameters():
        param.requires_grad = False
    for param in pipe.vae.parameters():
        param.requires_grad = False
    for param in pipe.text_encoder.parameters():
        param.requires_grad = False

    # 在UNet注意力层插入LoRA层
    lora_rank = 8
    lora_layers = []
    for name, module in pipe.unet.named_modules():
        if "attn" in name and isinstance(module, torch.nn.Linear):
            lora = LoRALinearLayer(module.in_features, module.out_features, rank=lora_rank).to(device, dtype=torch.float16)
            lora_layers.append(lora)
            setattr(module, "lora_layer", lora)

    # 加载图文数据集
    dataset = load_dataset(dataset_path, split="train")

    optimizer = torch.optim.AdamW(
        [{"params": lora_layers, "lr": 1e‑4}], weight_decay=1e‑4
    )

    # NPU开启AMP混合精度
    scaler = torch.npu.amp.GradScaler()
    num_epoch = 10
    pipe.unet.train()

    for epoch in range(num_epoch):
        for sample in dataset:
            pixel_values = sample["image"].to(device, dtype=torch.float16)
            prompt_ids = sample["prompt_ids"].to(device)

            with torch.npu.amp.autocast(dtype=torch.float16):
                noise = torch.randn_like(pixel_values)
                timestep = torch.randint(0, 1000, (pixel_values.shape[0],), device=device)
                model_pred = pipe.unet(pixel_values, timestep, encoder_hidden_states=prompt_ids).sample
                loss = torch.nn.functional.mse_loss(model_pred, noise)

            optimizer.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

        print(f"epoch {epoch}, loss: {loss.item():.4f}")

    # 保存LoRA权重safetensors格式
    os.makedirs(output_dir, exist_ok=True)
    pipe.save_lora_weights(save_directory=output_dir, unet_lora_layers=lora_layers)
    print(f"LoRA权重保存完成，路径：{output_dir}")

if __name__ == "__main__":
    main()

4.3 微调完成后推理验证脚本 infer_lora_npu.py

import torch
from diffusers import StableDiffusionPipeline

device = torch.device("npu:0")
base_model = "./stable‑diffusion‑v1‑5"
lora_path = "./sd_lora_output"

pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16).to(device)
pipe.load_lora_weights(lora_path)

prompt = "一只二次元风格小猫，油画风格"
image = pipe(prompt, num_inference_steps=25).images[0]
image.save("lora_gen_result.png")
print("生成图片已保存 lora_gen_result.png")

4.4 单机多卡分布式启动命令

# 4卡NPU分布式LoRA微调，底层使用HCCL通信
source env.sh
torchrun --nproc_per_node=4 train_lora_sd_npu.py

5 性能与显存优化效果参考

优化策略	效果说明
LoRA rank=8 + Gradient Checkpoint	显存占用降低约 70%，单卡可跑 SD1.5 微调
FP16 混合精度 AMP	显存下降 30‑40%，训练速度提升
DDP 4 卡分布式训练	训练时长缩短接近 3‑4 倍，依赖 HCCL 带宽
梯度累积 4 步	支持更小 batch，解决 OOM 显存溢出

提示：实际显存、训练速度取决昇腾硬件型号、图片分辨率、6 常见问题 FAQ

Q1：微调训练出现 NPU OOM 显存溢出怎么办？ A：1. 开启梯度检查点；2. 降低 batch_size；3. 开启梯度累积；4. 使用 LoRA 而不是全参数微调；5. 启用 FP16 混合精度。

Q2：部分算子回退 CPU 执行，训练速度很慢？ A：查看 PyTorch‑NPU 日志，确认 fallback 算子；升级 CANN、PyTorch‑NPU 版本；部分 Diffusers 自定义算子可替换为昇腾适配实现。

Q3：微调完成后生成图像效果不收敛、风格没有变化？ A：排查三点：数据集 prompt 质量；学习率、epoch 迭代次数超参；确认 LoRA 层正确挂载，主干网络权重被冻结。

Q4：微调后的 LoRA 权重如何部署到 MindIE 推理引擎？ A：把 LoRA 权重合并进基础 SD 模型，使用 ATC 工具转换为 om 离线模型，加载到 MindIE 完成图像生成推理服务部署。

Q5：多卡微调 HCCL 报错初始化失败？ A：确认 rank_table 配置正确，检查 HCCS/RoCE 链路，调大HCCL_CONNECT_TIMEOUT超时环境变量。

7 总结

图像生成 CV 模型微调是行业 AI 绘图落地关键环节，LoRA 轻量微调凭借低显存开销成为主流方案。昇腾平台基于 CANN、PyTorch‑NPU、HCCL 集合通信完整支持 Stable Diffusion 类扩散模型微调训练。

开发者通过梯度检查点、混合精度、LoRA、梯度累积等手段解决图像生成模型巨大显存压力；单机多卡场景使用 DDP+HCCL 分布式训练缩短训练周期；训练得到 LoRA 权重可以直接 Diffusers 验证，也可以合并后导出 om 模型部署 MindIE 推理服务。

在开发过程中结合 Profiling 工具观测 NPU AI Core 利用率，排查算子 CPU fallback，调优超参，就可以在昇腾 NPU 平台高效完成图像生成 CV 模型的定制化微调实践。