FLUX DanceGRPO 使用指南

目录

• 简介
• 环境安装
  • 仓库拉取
  • 环境搭建
• 权重下载
• 数据集准备及处理
• 训练
  • 准备工作
  • 启动训练
• 性能数据
• FAQ

简介

以 MindSpeed MM 仓库复现 DanceGRPO 后训练方法来帮助用户快速入门，前期需要完成代码仓、环境、数据集以及权重等准备工作，再按照说明中的启动方式启动训练，以下为具体的操作说明。

参考实现

DanceGRPO开源代码仓以及对应commit id如下：

url=https://github.com/XueZeyue/DanceGRPO
commit_id=2149f36f22db601f9dbf70472fea11576f62a0f6

环境安装

【模型开发时推荐使用配套的环境版本】

请参考安装指南

1. 仓库拉取

git clone https://gitcode.com/Ascend/MindSpeed-MM.git
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_v0.12.1
cp -r megatron ../MindSpeed-MM/
cd ..

cd MindSpeed-MM
mkdir -p logs data ckpt
cd ..

2. 环境搭建

请参考安装指南创建Python环境并安装 torch 和 TorchNPU。

conda activate test

# 安装加速库
git clone https://gitcode.com/Ascend/MindSpeed.git
cd MindSpeed
git checkout 5176c6f5f133111e55a404d82bd2dc14a809a6ab
cp -r mindspeed ../MindSpeed-MM/
cd ..

# 安装dance grpo依赖库
cd MindSpeed-MM
pip install -r ./examples/dancegrpo/requirements-lint.txt
cd ..

git clone https://github.com/tgxs002/HPSv2.git
cd HPSv2
git checkout 866735ecaae999fa714bd9edfa05aa2672669ee3
pip install -e .
cd ..

3.Decord搭建

【X86版安装】

pip install decord==0.6.0

【ARM版安装】

apt方式安装请参考链接

yum方式安装请参考脚本

权重下载

创建保存权重的目录：

cd MindSpeed-MM
mkdir ckpt/flux
mkdir ckpt/hps_ckpt
cd ..

下载FLUX预训练权重 FLUX预训练权重 ，下载至MindSpeed MM工程根目录下的ckpt/flux目录中。

下载HPS-v2.1预训练权重 HPS-v2.1预训练权重 ，将其中的HPS_v2.1_compressed.pt下载至MindSpeed MM工程根目录下的ckpt/hps_ckpt目录中。

下载CLIP预训练权重 CLIP预训练权重 ，将其中的open_clip_pytorch_model.bin下载至MindSpeed MM工程根目录下的ckpt/hps_ckpt目录中。

数据集准备及处理

下载FLUX DanceGRPO使用的提示词数据集。在文件页面点击download raw file下载文件至MindSpeed MM工程根目录的data目录下。

数据集下载完成后要对数据进行预处理，在启动预处理之前，可以根据自身训练配置需要修改数据预处理脚本的配置，以FLUX模型为例：

1. vae模型权重所在路径为LOAD_PATH，默认为ckpt/flux；
2. 预处理后的数据集存放路径为OUTPUT_DIR，默认为data/rl_embeddings；
3. 提示词文件路径为PROMPT_DIR，默认为data/prompts.txt。

上述注意点修改完毕后，可启动脚本进行数据预处理：

cd MindSpeed-MM
bash examples/dancegrpo/preprocess_flux_rl_embeddings.sh

处理后的数据默认会存储在MindSpeed MM根目录下的data/rl_embeddings目录中。

训练

1. 准备工作

配置脚本前需要完成前置准备工作，包括：环境安装、权重下载、数据集准备及处理，详情可查看对应章节。

2. 启动训练

以 FLUX 模型为例，在启动训练之前，可根据自身训练配置需要修改启动脚本的配置：

1. 根据使用机器的情况，修改NNODES、NPUS_PER_NODE配置， 例如单机8卡 可设置NNODES为 1 、NPUS_PER_NODE为8；
2. 如果为多机训练，需要保证各个节点的MASTER_ADDR一致，且为其中一台节点的IP；各节点的MASTER_PORT 配置为相同端口号；从IP为MASTER_ADDR的节点开始，将各节点的NODE_RANK配置为从0开始依次递增的整数；
3. 数据集配置信息路径为MM_DATA，默认路径为./examples/dancegrpo/data_dancegrpo.json；
4. 模型配置信息路径为MM_MODEL，默认路径为./examples/dancegrpo/model_dancegrpo.json；
5. DiT模型预训练权重加载路径为LOAD_PATH，默认路径为ckpt/flux，用户也可以根据自身权重存放位置进行调整；
6. 训练权重的保存路径为SAVE_PATH，默认为save_dir；
7. 模型训练过程的reward值保存文件的路径为HPS_REWARD_SAVE_PATH，默认为./hps_reward.txt。

在启动训练前，可根据自身训练配置需要修改数据集配置data_dancegrpo.json：

1. dataset_param.basic_parameters.data_path表示预处理数据中的元数据文件videos2caption.json的路径。

在启动训练前，可根据自身训练配置需要修改模型配置model_dancegrpo.json：

1. reward.ckpt_dir表示奖励模型预训练权重的路径。

上述注意点修改完毕后，可启动脚本开启训练：

bash examples/dancegrpo/posttrain_flux_dancegrpo.sh

注意：所有节点的代码、权重、数据等路径的层级要保持一致，且启动训练脚本的时候都位于MindSpeed MM目录下

训练完成后，会在logs目录中生成运行日志文件，生成训练reward记录文件。

Atlas 950 服务器性能测试脚本

如需在 Atlas 950 服务器上进行纯性能测试，可使用 posttrain_flux_dancegrpo_A5.sh 脚本。该脚本默认通过 --no-save_images 关闭采样图片的落盘，以减少 I/O 开销、专注于性能评估：

bash examples/dancegrpo/posttrain_flux_dancegrpo_A5.sh

注意：如需保留采样图片用于可视化，请将脚本中的 --no-save_images \ 一行删除，恢复默认的图片保存行为（图片会写入 ./images/flux_grpo/step{N}/ 目录下）。

---

性能数据

模型	机器型号	集群	任务	GBS	端到端 SPS
FLUX DanceGRPO	Atlas 200T A2 Box16	1*8	微调	32	0.1123

注：此处 SPS 代表 Samples per Second。

---

FAQ

1. 对于CPU型号为x86的设备，建议使用torchvision版本为0.25.0+cpu，若遇到有关torchvision的导包问题，建议优先检查环境中的torchvision版本是否为+cpu版本。