ColoDiffusion 实战指南：用 Colossal-AI 驯服 Stable Diffusion 训练、微调与推理

ColoDiffusion 实战指南：用 Colossal-AI 驯服 Stable Diffusion 训练、微调与推理

【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI

导读

本文围绕 Colossal-AI 官方示例ColoDiffusion（位于仓库examples/images/diffusion目录）展开，系统讲解如何借助 Colossal-AI 完成 Stable Diffusion v1/v2 从零训练、基于 DreamBooth / Teyvat 数据集的微调，以及 txt2img 文本生图推理。文中不仅完整继承示例文档的安装、训练与推理操作步骤，还会结合仓库内main.py、配置文件与启动脚本等源码，逐项解释关键参数的作用与底层调用链。读完本文，你将掌握一套可直接运行的 Stable Diffusion 训练流水线，并理解placement_policy（cuda / cpu / Gemini 异构内存）等 Colossal-AI 特性如何在大模型训练中发挥作用。

ColoDiffusion 是什么

ColoDiffusion 是 Colossal-AI 为 AIGC（AI-Generated Content）扩散模型加速而提供的参考实现，目标是让Stable Diffusion v1与Stable Diffusion v2这类模型的训练、微调与推理更省显存、更省钱、更快。官方 README（即 examples/images/diffusion/README.md）给出的典型收益包括：

• 训练：相比原生方案，Stable Diffusion 训练内存占用最多可降低约 5.6 倍，硬件成本最多可降低约 46 倍（从 A100 级别降至 RTX 3060 级别即可开展训练实验）；
• DreamBooth 微调：仅凭 3～5 张目标主体的图像即可个性化模型；
• 推理：推理阶段 GPU 显存占用可降低约 2.5 倍。

需要说明的是，上述数字来自示例 README 的自述结论，属于项目宣传口径；具体收益取决于模型规模、Batch Size、并行策略与实际硬件，应在自己的环境中实测验证。

该示例的 Roadmap 处于快速迭代状态，README 明确标注的状态为：

• 从零训练 Stable Diffusion v1 / v2；
• 微调预训练的 Stable Diffusion v1 模型；
• 使用 PyTorch 对预训练模型做推理；
• 微调预训练的 Stable Diffusion v2 模型；
• 使用 TensorRT 对预训练模型做推理（当时尚未完成）。

仓库目录速览

examples/images/diffusion/ ├── configs/ │ ├── train_colossalai.yaml # 全量预训练（Colossal-AI）配置 │ ├── train_ddp.yaml # DDP 对照训练配置 │ ├── train_colossalai_cifar10.yaml # CIFAR-10 小规模验证配置 │ ├── Inference/ # v2 推理专用配置（含 inpainting / 超分） │ └── Teyvat/ # Teyvat 角色微调配置与数据说明 ├── docker/Dockerfile # Docker 镜像构建文件 ├── environment.yaml # conda 环境（ldm）锁定版本 ├── main.py # 训练/微调统一入口 ├── train_colossalai.sh # Colossal-AI 启动脚本 ├── train_ddp.sh # PyTorch DDP 启动脚本 ├── ldm/ # latent diffusion 模型与数据层 └── scripts/txt2img.py # 文生图推理脚本

核心入口是 main.py，它基于 PyTorch Lightning 的Trainer封装训练流程，并通过strategy参数在ColossalAIStrategy与原生DDPStrategy之间切换。

环境安装（两种方式）

方式一：从源码 / pip 安装依赖

示例 README 建议使用 CUDA 11.6/11.8 的操作系统环境。仓库自带的 environment.yaml 已经锁定了该示例对应的完整依赖栈（conda环境名为ldm），可直接创建并激活：

conda env create -f environment.yaml conda activate ldm

该文件锁定的关键版本如下：

如果你已有 latent-diffusion 环境，也可以用如下命令就地补全依赖：

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch pip install transformers diffusers invisible-watermark

安装 Colossal-AI

README 明确指出：与本训练配套的推荐版本是colossalai==0.2.5（对应 torch 1.12.1 组合），既可以从官方源安装固定版本，也可以安装面向最新 torch 的最新版本：

# 推荐固定版本（与本次训练配套） pip install colossalai==0.2.5 # 面向最新 torch 版本 pip install colossalai

如需从源码安装 Colossal-AI，可先克隆本仓库再构建：

git clone https://gitcode.com/GitHub_Trending/co/ColossalAI cd ColossalAI BUILD_EXT=1 pip install .

需要留意的是，从源码结构看，当前 main.py 已改用新版 Lightning 的lightning.pytorch命名空间（from lightning.pytorch.strategies import ColossalAIStrategy, DDPStrategy），并在 scripts/txt2img.py 中为lightning.pytorch导入失败时保留了回退到pytorch_lightning的兼容分支。因此若你按environment.yaml中 lightning 1.9.x 的旧版本部署，遇到命名空间相关报错时需同步升级 Lightning 或调整导入，这一点在复现时请特别留意。

可选：xformers 加速

xformers中的 flash attention 实现可以加速训练，代价是额外的磁盘占用。示例 README 推荐的配套版本为 0.0.12：

pip install xformers==0.0.12

方式二：使用 Docker

示例提供了两条 Docker 路径：使用仓库内 docker/Dockerfile 自行构建，或直接拉取官方镜像：

# 1. 从 Dockerfile 构建 cd ColossalAI/examples/images/diffusion/docker docker build -t hpcaitech/diffusion:0.2.0 . # 2. 从 Docker Hub 拉取 docker pull hpcaitech/diffusion:0.2.0

启动容器与卷映射

镜像就绪后，建议在仓库根目录启动容器，并将当前目录挂载为工作区：

docker run --rm \ -it --gpus all \ -v $PWD:/workspace \ -v <your-data-dir>:/data/scratch \ -v <hf-cache-dir>:/root/.cache/huggingface \ hpcaitech/diffusion:0.2.0 \ /bin/bash

进入容器后即可按 README 流程训练：

cd examples/images/diffusion/ # 1. 先按下文步骤下载预训练 checkpoint # 2. 修改 train_colossalai.sh 中的配置 # 3. 启动 Colossal-AI 训练 bash train_colossalai.sh

Docker 卷映射有三个注意事项（README 原话要点）：

1. 必选：通过-v <your-data-dir>:/data/scratch把准备好的数据挂到/data/scratch。注意 Docker 内的路径需为 Linux 风格，例如 Windows 下的C:\User\Desktop应写成/mnt/c/User/Desktop；
2. 推荐：通过-v <hf-cache-dir>:/root/.cache/huggingface把模型权重缓存放宿主机，避免每次docker run重复下载预训练权重；
3. 可选：若容器内报共享内存不足，请追加-v /dev/shm:/dev/shm。

下载预训练 Checkpoint

训练需要先获取 Stable Diffusion 预训练权重，README 推荐了三个来源：

• stable-diffusion-v2-base（推荐）：使用 wget 从对应模型仓库直接下载权重文件512-base-ema.ckpt；
• stable-diffusion-v1-4（CompVis）：建议先执行git lfs install，再以 git lfs 方式克隆官方仓库stable-diffusion-v1-4；
• stable-diffusion-v1-5（runway）：同样以git lfs install后克隆仓库stable-diffusion-v1-5。

以上仓库均托管于 Hugging Face 的公开模型空间（模型标识分别为stabilityai/stable-diffusion-2-base、CompVis/stable-diffusion-v1-4、runwayml/stable-diffusion-v1-5），请前往对应页面确认最新直链后下载。下载得到的512-base-ema.ckpt正是后续训练命令中--ckpt参数所指的文件。

数据集准备

预训练阶段使用的数据集来自LAION-5B（LAION 组织的公开图文数据集子集）。你需要将数据集的实际路径写入训练配置中的data.file_path字段——即修改 configs/train_colossalai.yaml：

data: batch_size: 128 wrap: False # num_workers 建议为 2 * batch_size，且总数小于 1024 # 例如使用 8 个设备时，每个设备不超过 128 num_workers: 128 train: target: ldm.data.base.Txt2ImgIterableBaseDataset params: file_path: # 此处填写 YOUR DATASET_PATH world_size: 1 rank: 0

这里的数据集类型是ldm.data.base.Txt2ImgIterableBaseDataset（定义见 ldm/data/base.py），一个可迭代的图文数据集。world_size与rank用于在多进程下按 rank 切分数据。在 main.py 的worker_init_fn中可以看到：训练时每个 DataLoader worker 会把可迭代数据集按num_records // num_workers均分为互不重叠的片段，并为其设置独立随机种子，以避免多 worker 数据重复。

训练：Colossal-AI 版与 DDP 对照版

启动脚本

仓库同时提供了两个启动脚本，便于做性能对照：

• train_colossalai.sh：Colossal-AI 训练（主推方案）；
• train_ddp.sh：PyTorch DDP 训练。

两者的头部都设置了HF_DATASETS_OFFLINE=1 / TRANSFORMERS_OFFLINE=1 / DIFFUSERS_OFFLINE=1，强制 Hugging Face 相关库走离线模式（避免下载，配合已缓存权重使用）。

train_colossalai.sh中的核心命令为：

python main.py --logdir /tmp/ --train --base configs/train_colossalai.yaml --ckpt 512-base-ema.ckpt

对命令行参数的说明（README 要点）：

• --logdir：指定日志与最终 checkpoint 的保存目录。训练产物位置为：
  • checkpoint 位于logdir/checkpoints或logdir/diff_tb/version_0/checkpoints；
  • 训练配置 yaml 会被拷贝保存到logdir/configs；
• --ckpt：加载预训练权重（如512-base-ema.ckpt）；
• --base：指定用于构建模型的配置 yaml 路径。

main.py 的入口参数

除上面三个外，main.py 的解析器还支持一批常用参数，便于调试与恢复训练：

值得留意的是学习率缩放逻辑（main.py）：当--scale_lr开启时，实际学习率被设置为accumulate_grad_batches × 设备数 × batch_size × base_learning_rate；关闭则直接使用配置中的base_learning_rate。

训练配置逐项拆解

configs/train_colossalai.yaml 是完整的模型 + 数据 + Lightning 训练配置。README 对其中与 Colossal-AI 直接相关的训练参数给出如下约定：

• devices：训练使用的 GPU 数量，默认 8；
• max_epochs：最大训练轮数，默认 2；
• precision：训练精度类型，默认 16（即 fp16）。若要用 Colossal-AI 策略，必须使用 fp16；
• placement_policy：Colossal-AI 支持的显存策略，默认cuda（全部参数加载进 GPU 显存）；cpu表示 CPU offload 策略；auto则启用Gemini（异构内存管理器，同时利用 GPU 显存、CPU DRAM 与 NVMe SSD 存储参数）。

配置文件中的lightning.trainer段示例（来自该 yaml）：

lightning: trainer: accelerator: 'gpu' devices: 2 log_gpu_memory: all max_epochs: 2 precision: 16 auto_select_gpus: False strategy: use_chunk: True enable_distributed_storage: True placement_policy: cuda force_outputs_fp32: true min_chunk_size: 64 log_every_n_steps: 2 logger: True default_root_dir: "/tmp/diff_log/"

strategy段直接传入ColossalAIStrategy（在 main.py 中通过ColossalAIStrategy(**strategy_cfg)实例化；若配置里没有strategy段，则默认回退到DDPStrategy(find_unused_parameters=False)）。其中use_chunk（张量分块）、enable_distributed_storage（分布式存储）、force_outputs_fp32（输出强制 fp32）与min_chunk_size共同决定了 Gemini 对参数张量的分块与搬运粒度。关于 ColossalAIStrategy 的完整字段，可参考 PyTorch Lightning 官方对 model parallel 策略的文档。

模型侧的 UNet 配置也值得关注：unet_config.use_fp16: True与use_checkpoint: True配合 fp16 训练；context_dim: 1024对应 CLIP 文本编码器维度；attention_resolutions: [4, 2, 1]决定在哪些分辨率下启用注意力。scheduler_config中注释说明：从头训练时warm_up_steps应取 10000，而从已有 checkpoint 恢复时用 1。

main.py 训练主流程

从源码结构可以还原出 main.py 的完整执行链路，帮助理解配置如何生效：

1. 用OmegaConf.load依次加载--base指定的多个 yaml，再用命令行--key value（OmegaConf.from_dotlist）合并成最终配置；
2. 若命令行指定--ckpt，将其写入config.model.params.ckpt，随后实例化LatentDiffusion（定义见 ldm/models/diffusion/ddpm.py）；
3. 根据是否有lightning.logger配置决定使用 WandB 还是默认 TensorBoard logger（logdir/diff_tb）；
4. 按配置实例化ColossalAIStrategy或DDPStrategy；
5. 依次注册五个回调：
   • SetupCallback：创建 logdir/checkpoints/configs 目录，并把 project/lightning 配置以 yaml 形式归档到logdir/configs；同时监听键盘中断（on_keyboard_interrupt）时把last.ckpt存到 checkpoint 目录；
   • ImageLogger：按batch_frequency: 750、最多 4 张的频率把log_images的输出以网格（make_grid）形式保存到logdir/images/<split>，并把[-1,1]的像素值线性还原到[0,1]再转成 PNG（见 main.py）；
   • LearningRateMonitor：按 step 记录学习率曲线；
   • ModelCheckpoint：监控val/loss_simple_ema保存 top-3，同时每 10000 步保存trainstep_checkpoints权重；
   • CUDACallback：统计每个 epoch 的平均耗时与峰值显存（main.py）；
6. 数据侧由DataModuleFromConfig统一构造，其训练 DataLoader 使用自定义的DataLoaderX（继承自 DataLoader，迭代时套用BackgroundGenerator在后台预取数据，见 main.py）；
7. 最后注册SIGUSR1（存 checkpoint）与SIGUSR2（进入 pudb 调试）两个信号处理器，随后调用trainer.fit(model, data)。

由此可见，日志、可视化与断点续训全部由框架层封装，使用者只需维护一份 yaml 配置。

微调示例：Teyvat 数据集

除全量预训练外，示例还提供了在Teyvat数据集上的微调用例。该数据集由 BLIP 模型为《原神》角色图像生成文字描述（仓库内说明见 configs/Teyvat/README.md）：每行数据包含image与text两个键，image为尺寸不一的 PNG，text中串联了Teyvat、角色Name、Element、Weapon、Region、Model type等结构化标签以及 BLIP 生成的Description描述。

使用对应的微调配置启动训练（配置位于 configs/Teyvat/train_colossalai_teyvat.yaml）：

python main.py --logdir /tmp -t -b configs/Teyvat/train_colossalai_teyvat.yaml

与全量预训练配置相比，该配置主要有三处差异：

1. 数据目标不同：data.train.target为ldm.data.teyvat.hf_dataset，path: Fazzie/Teyvat直接使用 Hugging Face 上的 Teyvat 数据集（代码见 ldm/data/teyvat.py），无需像 LAION 那样准备本地文件；
2. 批量更小：batch_size: 16、num_workers: 4，更适合单机有限显存；
3. 在线图像增强：image_transforms依次执行Resize(512) → RandomCrop(512) → RandomHorizontalFlip，在训练时把图像统一裁剪到 512×512 并做随机水平翻转。

如果你希望做更细粒度的"主体定制"微调（用 3～5 张特定对象的图片个性化模型），仓库还单独提供了 DreamBooth 实现，说明文档见 examples/images/dreambooth/README.md，其中对--placement（cuda / cpu / auto-Gemini）等参数有与本文一致的语义解释，可作为进阶参考。

推理：txt2img 文本生图

直接测试预训练模型

README 给出的最基本推理命令如下（使用 PLMS 采样）：

python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms \ --outdir ./output \ --ckpt 512-base-ema.ckpt \ --config configs/train_ddp.yaml

其中--config用于指定构建模型结构的配置（既可以用训练配置，也可以用 configs/Inference/v2-inference.yaml 这类推理专用配置——后者只保留模型结构，不含lightning与data段）。针对 Stable Diffusion v2 生态，configs/Inference下还提供了 v2-inpainting、v2-midas（深度估计）、x4-upscaling（超分）等衍生任务的推理配置。

对自训练权重做推理

训练结束后，你会在--logdir下找到last.ckpt与归档的训练配置 yaml（路径形如logdir/configs/<时间戳>-project.yaml），可用它们替换推理命令中的权重与结构配置：

python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms \ --outdir ./output \ --ckpt path/to/logdir/checkpoints/last.ckpt \ --config path/to/logdir/configs/project.yaml

仓库自带的 scripts/txt2img.sh 也演示了这一用法——它用 Teyvat 微调产物生成角色图像：

python scripts/txt2img.py --prompt "Teyvat, Medium Female, a woman in a blue outfit holding a sword" --plms \ --outdir ./output \ --ckpt checkpoints/last.ckpt \ --config configs/2023-02-02T18-06-14-project.yaml \ --n_samples 4

txt2img.py 完整参数说明

scripts/txt2img.py 的参数定义与示例 README 中给出的 usage 一致，逐项整理如下：

结合 scripts/txt2img.py 的源码可见，脚本内部支持三种采样器：DDIMSampler、PLMSSampler与DPMSolverSampler（来自仓库内的 ldm/models/diffusion/dpm_solver 实现），并在加载权重后以torch.set_grad_enabled(False)关闭梯度、以autocast做 fp16 推理。README 中同时提到该推理路径可降低约 2.5 倍的 GPU 显存占用。除 txt2img 外，同目录还提供img2img.py、inpaint.py、sample_diffusion.py等脚本，可覆盖图生图、局部重绘等更多场景。

测试与工程化佐证

示例目录内还包含轻量测试脚本，可作为验证代码正确性的入口：

• scripts/tests/test_checkpoint.py：验证 checkpoint 的保存/加载是否一致；
• scripts/tests/test_watermark.py：验证输出图像的不可见水印；
• test_ci.sh：CI 入口脚本，可用于快速冒烟验证环境与主流程可运行。

引用与致谢

示例 README 明确致谢了其实现所依赖的开源工作，包括 OpenAI 的 guided-diffusion 代码库、lucidrains 的 denoising-diffusion-pytorch 与 x-transformers、Stable Diffusion、PyTorch Lightning 与 Hugging Face，以及 HazyResearch 的 flash attention 实现。相关研究可按以下 BibTeX 引用：

@article{bian2021colossal, title={Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training}, author={Bian, Zhengda and Liu, Hongxin and Wang, Boxiang and Huang, Haichen and Li, Yongbin and Wang, Chuanrui and Cui, Fan and You, Yang}, journal={arXiv preprint arXiv:2110.14883}, year={2021} } @misc{rombach2021highresolution, title={High-Resolution Image Synthesis with Latent Diffusion Models}, author={Robin Rombach and Andreas Blattmann and Dominik Lorenz and Patrick Esser and Björn Ommer}, year={2021}, eprint={2112.10752}, archivePrefix={arXiv}, primaryClass={cs.CV} } @article{dao2022flashattention, title={FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness}, author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher}, journal={arXiv preprint arXiv:2205.14135}, year={2022} }

小结

ColoDiffusion 是一套把 Colossal-AI 的异构显存能力（cuda / cpu offload / Gemini）接入 PyTorch Lightning + Stable Diffusion 生态的完整参考工程。通过它，你可以：

• 用一份 yaml 在 Colossal-AI 与 DDP 两种策略之间切换，量化对比训练吞吐与显存；
• 从 LAION-5B 规模数据开展预训练，或用 Teyvat、DreamBooth 等小数据做角色/主体微调；
• 用txt2img.py配合训练产物做端到端推理验证，支持 DDIM / PLMS / DPM-Solver 多种采样器。

实操前务必核对三件事：CUDA 版本与 torch/lightning 组合、Colossal-AI 版本与 fp16 前提（Colossal-AI 策略下 precision 必须为 16）、以及数据集在配置中的路径/离线缓存是否就绪。在此基础上，按train_colossalai.sh → 配置调优 → txt2img 推理的链路即可跑通整条 AIGC 训练流水线。

【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI