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_policycuda / cpu / Gemini 异构内存等 Colossal-AI 特性如何在大模型训练中发挥作用。ColoDiffusion 是什么ColoDiffusion 是 Colossal-AI 为 AIGCAI-Generated Content扩散模型加速而提供的参考实现目标是让Stable Diffusion v1与Stable Diffusion v2这类模型的训练、微调与推理更省显存、更省钱、更快。官方 README即 examples/images/diffusion/README.md给出的典型收益包括训练相比原生方案Stable Diffusion 训练内存占用最多可降低约 5.6 倍硬件成本最多可降低约 46 倍从 A100 级别降至 RTX 3060 级别即可开展训练实验DreamBooth 微调仅凭 35 张目标主体的图像即可个性化模型推理推理阶段 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该文件锁定的关键版本如下依赖版本python3.9.12cudatoolkit11.3pytorch / torchvision1.12.1 / 0.13.1colossalai0.2.5lightning1.9.0omegaconf2.1.1open_clip_torch2.0.2webdataset0.2.5如果你已有 latent-diffusion 环境也可以用如下命令就地补全依赖conda install pytorch1.12.1 torchvision0.13.1 torchaudio0.12.1 cudatoolkit11.3 -c pytorch pip install transformers diffusers invisible-watermark安装 Colossal-AIREADME 明确指出与本训练配套的推荐版本是colossalai0.2.5对应 torch 1.12.1 组合既可以从官方源安装固定版本也可以安装面向最新 torch 的最新版本# 推荐固定版本与本次训练配套 pip install colossalai0.2.5 # 面向最新 torch 版本 pip install colossalai如需从源码安装 Colossal-AI可先克隆本仓库再构建git clone https://gitcode.com/GitHub_Trending/co/ColossalAI cd ColossalAI BUILD_EXT1 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.12pip install xformers0.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.shDocker 卷映射有三个注意事项README 原话要点必选通过-v your-data-dir:/data/scratch把准备好的数据挂到/data/scratch。注意 Docker 内的路径需为 Linux 风格例如 Windows 下的C:\User\Desktop应写成/mnt/c/User/Desktop推荐通过-v hf-cache-dir:/root/.cache/huggingface把模型权重缓存放宿主机避免每次docker run重复下载预训练权重可选若容器内报共享内存不足请追加-v /dev/shm:/dev/shm。下载预训练 Checkpoint训练需要先获取 Stable Diffusion 预训练权重README 推荐了三个来源stable-diffusion-v2-base推荐使用 wget 从对应模型仓库直接下载权重文件512-base-ema.ckptstable-diffusion-v1-4CompVis建议先执行git lfs install再以 git lfs 方式克隆官方仓库stable-diffusion-v1-4stable-diffusion-v1-5runway同样以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-5BLAION 组织的公开图文数据集子集。你需要将数据集的实际路径写入训练配置中的data.file_path字段——即修改 configs/train_colossalai.yamldata: 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.shColossal-AI 训练主推方案train_ddp.shPyTorch DDP 训练。两者的头部都设置了HF_DATASETS_OFFLINE1 / TRANSFORMERS_OFFLINE1 / DIFFUSERS_OFFLINE1强制 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 的解析器还支持一批常用参数便于调试与恢复训练参数默认值说明-n / --namelogdir 的后缀名-r / --resume从 logdir 或其中 checkpoint 恢复训练与--name互斥-b / --base空列表基础配置从左到右依次加载并合并支持用--key value覆盖-t / --trainFalse是否执行训练-p / --project—新建或已有 project 名称-c / --ckpt加载 Stable AI 预训练 checkpoint-d / --debugFalse出错时进入事后调试pudb/pdb-s / --seed23随机种子-f / --postfix默认名称的附加后缀-l / --logdirlogs日志目录--scale_lrTrue按ngpu * batch_size * n_accumulate缩放基础学习率值得留意的是学习率缩放逻辑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 数量默认 8max_epochs最大训练轮数默认 2precision训练精度类型默认 16即 fp16。若要用 Colossal-AI 策略必须使用 fp16placement_policyColossal-AI 支持的显存策略默认cuda全部参数加载进 GPU 显存cpu表示 CPU offload 策略auto则启用Gemini异构内存管理器同时利用 GPU 显存、CPU DRAM 与 NVMe SSD 存储参数。配置文件中的lightning.trainer段示例来自该 yamllightning: 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_parametersFalse)。其中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 的完整执行链路帮助理解配置如何生效用OmegaConf.load依次加载--base指定的多个 yaml再用命令行--key valueOmegaConf.from_dotlist合并成最终配置若命令行指定--ckpt将其写入config.model.params.ckpt随后实例化LatentDiffusion定义见 ldm/models/diffusion/ddpm.py根据是否有lightning.logger配置决定使用 WandB 还是默认 TensorBoard loggerlogdir/diff_tb按配置实例化ColossalAIStrategy或DDPStrategy依次注册五个回调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.pyLearningRateMonitor按 step 记录学习率曲线ModelCheckpoint监控val/loss_simple_ema保存 top-3同时每 10000 步保存trainstep_checkpoints权重CUDACallback统计每个 epoch 的平均耗时与峰值显存main.py数据侧由DataModuleFromConfig统一构造其训练 DataLoader 使用自定义的DataLoaderX继承自 DataLoader迭代时套用BackgroundGenerator在后台预取数据见 main.py最后注册SIGUSR1存 checkpoint与SIGUSR2进入 pudb 调试两个信号处理器随后调用trainer.fit(model, data)。由此可见日志、可视化与断点续训全部由框架层封装使用者只需维护一份 yaml 配置。微调示例Teyvat 数据集除全量预训练外示例还提供了在Teyvat数据集上的微调用例。该数据集由 BLIP 模型为《原神》角色图像生成文字描述仓库内说明见 configs/Teyvat/README.md每行数据包含image与text两个键image为尺寸不一的 PNGtext中串联了Teyvat、角色Name、Element、Weapon、Region、Model type等结构化标签以及 BLIP 生成的Description描述。使用对应的微调配置启动训练配置位于 configs/Teyvat/train_colossalai_teyvat.yamlpython main.py --logdir /tmp -t -b configs/Teyvat/train_colossalai_teyvat.yaml与全量预训练配置相比该配置主要有三处差异数据目标不同data.train.target为ldm.data.teyvat.hf_datasetpath: Fazzie/Teyvat直接使用 Hugging Face 上的 Teyvat 数据集代码见 ldm/data/teyvat.py无需像 LAION 那样准备本地文件批量更小batch_size: 16、num_workers: 4更适合单机有限显存在线图像增强image_transforms依次执行Resize(512) → RandomCrop(512) → RandomHorizontalFlip在训练时把图像统一裁剪到 512×512 并做随机水平翻转。如果你希望做更细粒度的主体定制微调用 35 张特定对象的图片个性化模型仓库还单独提供了 DreamBooth 实现说明文档见 examples/images/dreambooth/README.md其中对--placementcuda / 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 4txt2img.py 完整参数说明scripts/txt2img.py 的参数定义与示例 README 中给出的 usage 一致逐项整理如下参数说明-h / --help显示帮助--prompt [PROMPT]要渲染的文本提示--outdir [OUTDIR]结果输出目录--skip_grid不保存网格图、只保存单张样本评估大量样本时有用--skip_save不保存单张样本仅用于测速--ddim_steps DDIM_STEPSDDIM 采样步数--plms使用 PLMS 采样--laion400m使用 LAION400M 训练出的模型--fixed_code所有样本使用同一个起始噪声便于对照--ddim_eta DDIM_ETADDIM etaeta0.0 对应确定性采样--n_iter N_ITER重复采样的轮数--H H输出图像高度像素空间--W W输出图像宽度像素空间--C Clatent 通道数--f F下采样倍率--n_samples N_SAMPLES每个 prompt 生成的样本数即 batch size--n_rows N_ROWS网格中的行数默认等于 n_samples--scale SCALE无分类器引导强度eps eps(x, empty) scale * (eps(x, cond) - eps(x, empty))--from-file FROM_FILE从文件逐行读取 prompt--config CONFIG构建模型的配置路径--ckpt CKPT模型 checkpoint 路径--seed SEED随机种子保证采样可复现--use_int8是否启用量化--precision {full,autocast}评估时使用的精度结合 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.shCI 入口脚本可用于快速冒烟验证环境与主流程可运行。引用与致谢示例 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考