scripts/train_ddpo_pairs_multigpu.sh
#!/bin/bash
#SBATCH --job-name=ddpo                     # Название задачи
#SBATCH --output=ddpo_%j.out                # stdout
#SBATCH --error=ddpo_%j.err                 # stderr
#SBATCH --time=3-00:00:00                   # Максимальное время выполнения
#SBATCH --nodes 1                           # Все CPU и GPU будут на одном вычислительном узле
#SBATCH --gpus 1                            # Требуемое кол-во GPU
#SBATCH --cpus-per-task 4                   # Количество выделенных ядер CPU

module restore
source .venv/bin/activate

which python
python -V
nvidia-smi


# Effective BS will be (N_GPU * train_batch_size * gradient_accumulation_steps)
# Paper used 2048. Training takes ~24 hours / 2000 steps
# bs=4 for v100, bs=8 for a100

export WANDB_MODE="offline"
accelerate launch \
  --multi_gpu \
  --num_processes 2 \
  --num_machines 1 \
  train_ddpo.py \
  --pretrained_model_name_or_path experiments/dreambooth/00000-3417-dog6-lr2 \
  --dataset_path data/pairs-TS-100k.json \
  --output_dir experiments/ddpo/00002-4d11-dog6-TS \
  --beta_dpo 5000 \
  --max_train_steps=2000 \
  --train_batch_size=2 \
  --gradient_accumulation_steps=64 \
  --learning_rate=1e-8 --scale_lr \
  --lr_scheduler="constant_with_warmup" --lr_warmup_steps=500 \
  --checkpointing_steps 200 \
  --dataloader_num_workers=4 \
  --hard_skip_resume \
  --report_to wandb \
  --tracker_project_name persgen \
  --validation_steps 10 \
  --validation_concept "sks dog" \
  --validation_concept_image_dir dreambooth/dataset/dog6 \
  --val_batch_size 4
