ByteDance-Seed/Bagel项目训练全流程指南：从数据准备到模型调优

项目概述

ByteDance-Seed/Bagel是一个多模态AI训练框架，支持文本到图像生成(T2I)、图像编辑(Editing)和视觉语言模型(VLM)等多种任务。本文将详细介绍该项目的完整训练流程，包括数据准备、训练配置和模型调优等关键环节。

数据准备详解

数据集来源

项目提供了三种典型任务的示例数据集：

1. 文本到图像(T2I)：基于FLUX.1-dev生成的高质量图像-文本对
2. 图像编辑(Editing)：来自SEED-Data-Edit-Part3的真实编辑样本
3. 视觉语言模型(VLM)：源自LLaVA-OneVision-Data的视觉-语言对话数据

数据格式支持

项目支持两种主要数据格式：

• 原始图像文件夹：传统目录结构存储的JPEG/PNG图像
• Parquet分片：高效的列式存储格式，适合大规模数据集

数据准备步骤

1. 下载示例数据集

   wget -O bagel_example.zip \
     https://lf3-static.bytednsdoc.com/obj/eden-cn/nuhojubrps/bagel_example.zip
   unzip bagel_example.zip -d /data
   

2. 目录结构说明

   bagel_example
   ├── t2i/                    # 文本到图像任务(Parquet格式)
   ├── editing/                # 图像编辑任务
   │   ├── seedxedit_multi/    # 多条件编辑样本
   │   └── parquet_info/       # Parquet元信息
   └── vlm/                    # 视觉语言模型
       ├── images/             # 原始图像(JPEG/PNG)
       └── llava_ov_si.jsonl   # 视觉-语言对话数据
   

3. 配置数据集路径 修改data/dataset_info.py文件中的所有your_data_path占位符，指向实际数据路径。

4. 扩展自定义数据集(可选) 可以通过扩展DATASET_INFO字典来添加自定义的Parquet分片或JSONL文件。

训练流程详解

基础训练命令

项目提供两种主要训练模式：

1. 预训练模式

torchrun \
  --nnodes=$num_nodes \
  --node_rank=$node_rank \
  --nproc_per_node=8 \
  --master_addr=$master_addr \
  --master_port=$master_port \
  train/pretrain_unified_navit.py \
  --dataset_config_file ./data/configs/example.yaml \
  --llm_path $llm_path \
  --vae_path $vae_path \
  --vit_path $vit_path \
  --layer_module Qwen2MoTDecoderLayer \
  --use_flex True \
  --resume_from $resume_from \
  --results_dir $output_path \
  --checkpoint_dir $ckpt_path \
  --max_latent_size 64  # 低分辨率预训练可设为32

2. 微调模式

torchrun \
  --nnodes=$num_nodes \
  --node_rank=$node_rank \
  --nproc_per_node=8 \
  --master_addr=$master_addr \
  --master_port=$master_port \
  train/pretrain_unified_navit.py \
  --dataset_config_file ./data/configs/example.yaml \
  --model_path $model_path \
  --layer_module Qwen2MoTDecoderLayer \
  --max_latent_size 64 \
  --resume-from $model_path \
  --finetune_from_hf True \
  --auto_resume True \
  --resume-model-only True \
  --finetune-from-ema True \
  --log_every 1 \
  --lr 2e-5 \
  --num_worker 1 \
  --expected_num_tokens 10240 \
  --max_num_tokens 11520 \
  --max_num_tokens_per_sample 10240

关键训练参数说明

模型配置参数

参数	默认值	说明
llm_path	hf/Qwen2.5-0.5B-Instruct	语言模型主干网络
vae_path	flux/vae/ae.safetensors	预训练的VAE检查点
vit_path	hf/siglip-so400m-14-980-flash-attn2-navit	视觉理解用的SigLIP ViT
max_latent_size	32	最大潜在网格尺寸，决定生成分辨率
text_cond_dropout_prob	0.1	文本条件丢弃概率

数据配置参数

参数	默认值	说明
dataset_config_file	data/configs/example.yaml	数据集分组和采样权重配置文件
max_num_tokens_per_sample	16384	单个样本最大token数限制
max_num_tokens	36864	批量处理时的最大token数上限

训练优化参数

参数	默认值	说明
total_steps	500_000	优化器总步数
lr	1e-4	峰值学习率
ema	0.9999	模型权重的指数移动平均衰减率
use_flex	True	启用FLEX打包提高GPU利用率

训练技巧与注意事项

1. 分辨率设置：

   • 微调时必须设置max_latent_size=64以确保加载正确的预训练权重
   • 低分辨率预训练可使用max_latent_size=32

2. 数据工作器配置：

   • num_used_data总值应大于NUM_GPUS × NUM_WORKERS
   • 小规模数据测试时可设num_worker=1

3. 任务特定调整：

   • 纯T2I微调：设置visual_und=False
   • 纯VLM微调：设置visual_gen=False

4. 调试建议：

   • 可减小expected_num_tokens等参数值加速调试
   • 典型调试输出示例：

     [时间戳] (step=0000000) Train Loss mse: 0.4063, Train Loss ce: 0.5504
     [时间戳] (step=0000001) Train Loss mse: 0.4121, Train Loss ce: 0.8152
     

5. 分布式训练：

   • 多节点训练需正确设置num_nodes和node_rank
   • 每个节点的GPU数量通过nproc_per_node指定

模型调优建议

1. 学习率策略：

   • 预训练阶段可采用constant或cosine调度
   • 微调阶段建议使用较小的学习率(如2e-5)

2. 正则化技巧：

   • 适当调整各种dropout概率防止过拟合
   • 使用梯度裁剪(max_grad_norm)稳定训练

3. 内存优化：

   • 大模型可冻结部分模块(freeze_llm等)节省显存
   • 调整max_num_tokens等参数控制显存占用

4. 日志与监控：

   • 配置W&B进行训练可视化
   • 合理设置log_every平衡监控频率与性能

通过灵活调整上述参数，开发者可以根据自身硬件条件和数据规模，优化训练流程，获得最佳模型性能。