重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新的库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

重要提示

在开始本教程之前，请务必查看简介，了解有关设置 NeMo-Aligner 环境的提示。

如果您遇到任何问题，请参阅 NeMo 的已知问题页面。该页面列举了已知问题，并在适当情况下提供了建议的解决方法。

完成本教程后，请参阅评估文档，了解有关评估训练模型的提示。

使用 DRaFT+ 微调 Stable Diffusion

在本教程中，我们将逐步介绍如何使用 NVIDIA 的 DRaFT+ 算法微调 Stable Diffusion 模型。DRaFT+ 通过正则化来缓解模式崩溃并提高多样性，从而增强了 DRaFT DRaFT 算法。有关 DRaFT+ 算法的更多技术细节，请查看我们的技术博客。

运行 DRaFT+ 的数据输入

运行 DRaFT+ 的数据应为 .tar 文件，其中包含纯文本提示。您可以从 .txt 文件生成 tar 文件，其中包含以换行符分隔的提示，格式如下

prompt1
prompt2
prompt3
prompt4
...

使用以下脚本从 Pick a pic 数据集下载并保存提示

from datasets import load_dataset

dataset = load_dataset("yuvalkirstain/pickapic_v1_no_images")
captions = dataset['train']['caption']
file_path = # path to save as a .txt file
with open(file_path, 'w') as file:
    for caption in captions:
        file.write(caption + '\n')

然后，您可以运行以下代码片段将其转换为 .tar 文件

import webdataset as wds

txt_file_path = # Path for the input txt file
tar_file_name = # Path for the output tar file

with open(txt_file_path, 'r') as f:
    prompts = f.readlines()
prompts = [item.strip() for item in prompts]
sink = wds.TarWriter(tar_file_name)
for index, line in enumerate(prompts):
    sink.write({
        "__key__": "sample%06d" % index,
        "txt": line.strip(),
    })
sink.close()

奖励模型

目前，我们仅支持 Pickscore 风格的奖励模型 (PickScore/HPSv2)。由于 Pickscore 是一个基于 CLIP 的模型，您可以使用 NeMo 的转换脚本将其从 huggingface 转换为 NeMo。

DRaFT+ 训练

要开始奖励模型训练，您需要训练好的 Stable Diffusion 模型的 UNet 和 VAE 组件的检查点，以及奖励模型的检查点。

终端

要直接在终端上运行 DRaFT+

GPFS="/path/to/nemo-aligner-repo"
TRAIN_DATA_PATH="/path/to/train_dataset.tar"
UNET_CKPT="/path/to/unet_weights.ckpt"
VAE_CKPT="/path/to/vae_weights.bin"
RM_CKPT="/path/to/reward_model.nemo"
DRAFTP_SCRIPT="train_sd_draftp.py"       # or train_sdxl_draftp.py

torchrun --nproc_per_node=2 ${GPFS}/examples/mm/stable_diffusion/${DRAFTP_SCRIPT} \
   trainer.num_nodes=1 \
   trainer.devices=2 \
   model.micro_batch_size=1 \
   model.global_batch_size=8 \
   model.kl_coeff=0.2 \
   model.optim.lr=0.0001 \
   model.unet_config.from_pretrained=${UNET_CKPT} \
   model.first_stage_config.from_pretrained=${VAE_CKPT} \
   rm.model.restore_from_path=${RM_CKPT} \
   model.data.train.webdataset.local_root_path=${TRAIN_DATA_PATH} \
   exp_manager.create_wandb_logger=False \
   exp_manager.explicit_log_dir=/results

Slurm

要使用 Slurm 运行 DRaFT+。以下脚本使用 1 个节点

#!/bin/bash
#SBATCH -A <<ACCOUNT NAME>>
#SBATCH -p <<PARTITION NAME>>
#SBATCH -N 4
#SBATCH -t 4:00:00
#SBATCH -J <<JOB NAME>>
#SBATCH --ntasks-per-node=8
#SBATCH --exclusive
#SBATCH --overcommit

GPFS="/path/to/nemo-aligner-repo"

GPFS="/path/to/nemo-aligner-repo"
TRAIN_DATA_PATH="/path/to/train_dataset.tar"
UNET_CKPT="/path/to/unet_weights.ckpt"
VAE_CKPT="/path/to/vae_weights.bin"
RM_CKPT="/path/to/reward_model.nemo"

PROJECT="<<WANDB PROJECT>>"

CONTAINER=<<<CONTAINER>>> # use the latest NeMo Training container, Aligner will work there

MOUNTS="--container-mounts=MOUNTS" # mounts

RESULTS_DIR="/path/to/result_dir"

OUTFILE="${RESULTS_DIR}/rm-%j_%t.out"
ERRFILE="${RESULTS_DIR}/rm-%j_%t.err"
mkdir -p ${RESULTS_DIR}

MOUNTS="--container-mounts=MOUNTS" # mounts

DRAFTP_SCRIPT="train_sd_draftp.py"       # or train_sdxl_draftp.py

read -r -d '' cmd <<EOF
echo "*******STARTING********" \
&& echo "---------------" \
&& echo "Starting training" \
&& cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/mm/stable_diffusion/${DRAFTP_SCRIPT} \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   model.micro_batch_size=2 \
   model.global_batch_size=16 \
   model.kl_coeff=0.2 \
   model.optim.lr=0.0001 \
   model.unet_config.from_pretrained=${UNET_CKPT} \
   model.first_stage_config.from_pretrained=${VAE_CKPT} \
   rm.model.restore_from_path=${RM_CKPT} \
   model.data.webdataset.local_root_path=${TRAIN_DATA_PATH} \
   exp_manager.explicit_log_dir=${RESULTS_DIR} \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=${NAME} \
   exp_manager.wandb_logger_kwargs.project=${PROJECT}
EOF

srun --no-container-mount-home -o $OUTFILE -e $ERRFILE --container-image=$CONTAINER $MOUNTS bash -c "${cmd}"
set +x

注意

有关 DRaFT+ 超参数的更多信息，请参阅模型配置文件（分别为 SD 和 SDXL）

NeMo-Aligner/examples/mm/stable_diffusion/conf/draftp_sd.yaml NeMo-Aligner/examples/mm/stable_diffusion/conf/draftp_sdxl.yaml

DRaFT+ 结果

完成使用 DRaFT+ 微调 Stable Diffusion 后，您可以使用 NeMo 代码库中的 sd_infer.py 和 sd_lora_infer.py 脚本对已保存的模型运行推理。使用微调模型生成的图像应具有更好的提示对齐和美学质量。

使用退火重要性采样 (AIG) 的用户可控微调

AIG 提供了推理时的灵活性，可以在基础 Stable Diffusion 模型（低奖励和高多样性）和 DRaFT+ 微调模型（高奖励和低多样性）之间进行插值，以获得具有高奖励和高多样性的图像。通过指定逗号分隔的 weight_type 策略在基础模型和微调模型之间进行插值，可以轻松完成 AIG 推理。

Stable Diffusion XL 上的 AIG

weight type 为 base 时，使用基础模型进行 AIG；draft 使用微调模型（两种情况下都不进行插值）。weight type 为 power_<float> 形式时，使用 AIG 论文中指定的指数衰减进行插值。

要直接在终端上运行 AIG 推理

NUMNODES=1
LR=${LR:=0.00025}
INF_STEPS=${INF_STEPS:=25}
KL_COEF=${KL_COEF:=0.1}
ETA=${ETA:=0.0}
DATASET=${DATASET:="pickapic50k.tar"}
MICRO_BS=${MICRO_BS:=1}
GRAD_ACCUMULATION=${GRAD_ACCUMULATION:=4}
PEFT=${PEFT:="sdlora"}
NUM_DEVICES=${NUM_DEVICES:=8}
GLOBAL_BATCH_SIZE=$((MICRO_BS*NUM_DEVICES*GRAD_ACCUMULATION*NUMNODES))
LOG_WANDB=${LOG_WANDB:="False"}

echo "additional kwargs: ${ADDITIONAL_KWARGS}"

WANDB_NAME=SDXL_Draft_annealing
WEBDATASET_PATH=/path/to/${DATASET}

CONFIG_PATH="/opt/nemo-aligner/examples/mm/stable_diffusion/conf"
CONFIG_NAME=${CONFIG_NAME:="draftp_sdxl"}
UNET_CKPT="/path/to/unet.ckpt"
VAE_CKPT="/path/to/vae.ckpt"
RM_CKPT="/path/to/reward_model.nemo"
PROMPT=${PROMPT:="Bananas growing on an apple tree"}
DIR_SAVE_CKPT_PATH=/path/to/explicit_log_dir

if [ ! -z "${ACT_CKPT}" ]; then
    ACT_CKPT="model.activation_checkpointing=$ACT_CKPT "
    echo $ACT_CKPT
fi

EVAL_SCRIPT=${EVAL_SCRIPT:-"anneal_sdxl.py"}
export DEVICE="0,1,2,3,4,5,6,7" && echo "Running DRaFT+ on ${DEVICE}" && export HYDRA_FULL_ERROR=1
set -x
CUDA_VISIBLE_DEVICES="${DEVICE}" torchrun --nproc_per_node=$NUM_DEVICES /opt/nemo-aligner/examples/mm/stable_diffusion/${EVAL_SCRIPT} \
    --config-path=${CONFIG_PATH} \
    --config-name=${CONFIG_NAME} \
    model.optim.lr=${LR} \
    model.optim.weight_decay=0.0005 \
    model.optim.sched.warmup_steps=0 \
    model.sampling.base.steps=${INF_STEPS} \
    model.kl_coeff=${KL_COEF} \
    model.truncation_steps=1 \
    trainer.draftp_sd.max_epochs=5 \
    trainer.draftp_sd.max_steps=10000 \
    trainer.draftp_sd.save_interval=200 \
    trainer.draftp_sd.val_check_interval=20 \
    trainer.draftp_sd.gradient_clip_val=10.0 \
    model.micro_batch_size=${MICRO_BS} \
    model.global_batch_size=${GLOBAL_BATCH_SIZE} \
    model.peft.peft_scheme=${PEFT} \
    model.data.webdataset.local_root_path=$WEBDATASET_PATH \
    rm.model.restore_from_path=${RM_CKPT} \
    trainer.devices=${NUM_DEVICES} \
    trainer.num_nodes=${NUMNODES} \
    rm.trainer.devices=${NUM_DEVICES} \
    rm.trainer.num_nodes=${NUMNODES} \
    +prompt="${PROMPT}" \
    exp_manager.create_wandb_logger=${LOG_WANDB} \
    model.first_stage_config.from_pretrained=${VAE_CKPT} \
    model.first_stage_config.from_NeMo=True \
    model.unet_config.from_pretrained=${UNET_CKPT} \
    model.unet_config.from_NeMo=True \
    $ACT_CKPT \
    exp_manager.wandb_logger_kwargs.name=${WANDB_NAME} \
    exp_manager.resume_if_exists=True \
    exp_manager.explicit_log_dir=${DIR_SAVE_CKPT_PATH} \
    exp_manager.wandb_logger_kwargs.project=${PROJECT} +weight_type='draft,base,power_2.0'