写在最前面:

全量微调qwen-image模型需要大量显存,亲测使用小训练集1个batch仍然需要120+G VRAM。因此如果没有足够显存不要尝试训练qwen-image底模,可以训练LoRA或者可以尝试其他训练项目和方法。输出的每个模型约有70+G,因此也需要有足够的磁盘空间以保障模型正常训练。

一、训练项目

https://github.com/kohya-ss/musubi-tuner

本项目使用 musubi-tuner 项目进行底模训练,先clone项目再参考readme进行环境配置。

二、数据集准备

1.新建.toml文件进行训练配置

训练前需要新建一个.toml文件保存配置,可以参考如下模板。

# resolution, caption_extension, batch_size, num_repeats, enable_bucket, bucket_no_upscale should be set in either general or datasets
# otherwise, the default values will be used for each item

# general configurations
[general]
resolution = [960, 544]
caption_extension = ".txt"
batch_size = 1
enable_bucket = true
bucket_no_upscale = false

[[datasets]]
image_directory = "/path/to/image_dir"
cache_directory = "/path/to/cache_directory"
num_repeats = 1 # optional, default is 1. Number of times to repeat the dataset. Useful to balance the multiple datasets with different sizes.

# other datasets can be added here. each dataset can have different configurations

cache_directory是可选项,默认值为 None,表示使用与图像目录相同的目录。不过,作者建议设置缓存目录,以避免不同数据集之间意外共享缓存文件。

image_directory的格式是:训练图片+存有对应prompt同名的txt文件

使用.txt文档保存图片tag时可以参考上面的模板进行配置,使用json文件以及其他更详细的信息可以参考https://github.com/kohya-ss/musubi-tuner/blob/main/docs/dataset_config.md 文档。
 

2.Pre-caching

使用该项目训练前需要先进行 Latent Pre-caching 和 Text Encoder Output Pre-caching。

Latent Pre-caching:

python src/musubi_tuner/qwen_image_cache_latents.py \
    --dataset_config path/to/toml \
    --vae path/to/vae_model

Text Encoder Output Pre-caching:

python src/musubi_tuner/qwen_image_cache_text_encoder_outputs.py \
    --dataset_config path/to/toml \
    --text_encoder path/to/text_encoder \
    --batch_size 1

三、训练模型

项目使用 qwen_image_train.py 脚本对qwen-image模型进行全量微调,参数参考如下。

accelerate launch --num_cpu_threads_per_process 1 src/musubi_tuner/qwen_image_train.py \
    --dit path/to/dit_model \
    --vae path/to/vae_model \
    --text_encoder path/to/text_encoder \
    --dataset_config path/to/toml \
    --sdpa --mixed_precision bf16 --gradient_checkpointing \
    --optimizer_type adafactor --learning_rate 1e-6 --fused_backward_pass \
    --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" \
    --max_grad_norm 0 --lr_scheduler constant_with_warmup --lr_warmup_steps 10 \
    --max_data_loader_n_workers 2 --persistent_data_loader_workers \
    --max_train_epochs 16 --save_every_n_epochs 1 --seed 42 \
    --output_dir path/to/output_dir --output_name name-of-model

输入diffusers格式的dit模型时,--dit的路径配置为模型的transformer子文件夹中的第一个分片模型,文件名一般为 transformer/diffusion_model-00001-of-0000x.safetensors 

如需在训练过程中输出sample,可以加入命令 --sample_at_first --sample_prompts path/to/prompt.txt  注意加入采样会导致更高的VRAM消耗,很可能会OOM。

参数说明:

微调需要大量的显存。强烈建议使用内存节省选项。
--full_bf16:以 bfloat16 格式加载模型权重,以显著减少显存使用量。
--optimizer_type adafactor:对于微调,建议使用 Adafactor。
--fused_backward_pass:使用 Adafactor 时,在反向传播过程中减少显存使用量。
--mem_eff_save:在保存检查点时减少主内存(RAM)使用量。
--blocks_to_swap:在显存和主内存之间交换模型块以减少显存使用量。当显存有限时,这种方法有效。
--disable_numpy_memmap:禁用模型加载时的 numpy 内存映射,使用标准文件读取方式进行加载。会增加内存使用量,但在某些情况下可能会加快模型加载速度。
--full_bf16 可减少约 20GB 的显存使用量,但可能会影响模型精度,因为权重仍以 bfloat16 格式保存。请注意,优化器状态仍保持在 float32 中。此外,建议与支持随机舍入的优化器一起使用此选项。在本存储库中,带有 --fused_backward_pass 选项的 Adafactor 优化器支持随机舍入。

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐