[feat] Add Wan RL DiffusionNFT training (#1450)
Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
This commit is contained in:
Junda Su
co-authored by
Peiyuan Zhang
Peiyuan Zhang
parent
30e45c2411
commit
5854aec2ce
@@ -0,0 +1,180 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Build a text-only Parquet dataset from a one-prompt-per-line file. The script
|
||||
# shards prompts across GPU_NUM single-GPU torchrun workers, runs
|
||||
# v1_preprocess.py with --preprocess_task text_only, and writes prompt
|
||||
# embeddings/captions under OUTPUT_DIR for DMD2/DiffusionNFT text-only runs.
|
||||
|
||||
INPUT_FILE="${1:-train.txt}"
|
||||
OUTPUT_DIR="${OUTPUT_DIR:-data/train_text_only_dmd_preprocessed}"
|
||||
MODEL_PATH="${MODEL_PATH:-Wan-AI/Wan2.1-T2V-1.3B-Diffusers}"
|
||||
GPU_NUM="${GPU_NUM:-2}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-1}"
|
||||
SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-8}"
|
||||
FLUSH_FREQUENCY="${FLUSH_FREQUENCY:-8}"
|
||||
TEXT_MAX_LENGTH="${TEXT_MAX_LENGTH:-512}"
|
||||
CONDA_ROOT="${CONDA_ROOT:-/root/miniconda3}"
|
||||
CONDA_ENV="${CONDA_ENV:-fastvideo}"
|
||||
MIN_FREE_GPU_MB="${MIN_FREE_GPU_MB:-22000}"
|
||||
|
||||
if [[ ! -f "$INPUT_FILE" ]]; then
|
||||
echo "Input text file not found: $INPUT_FILE" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ ! -f "$CONDA_ROOT/etc/profile.d/conda.sh" ]]; then
|
||||
echo "Conda activation script not found under $CONDA_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# shellcheck source=/dev/null
|
||||
source "$CONDA_ROOT/etc/profile.d/conda.sh"
|
||||
conda activate "$CONDA_ENV"
|
||||
|
||||
if [[ "${HF_HUB_ENABLE_HF_TRANSFER:-0}" == "1" ]]; then
|
||||
if ! python -c "import hf_transfer" >/dev/null 2>&1; then
|
||||
echo "HF_HUB_ENABLE_HF_TRANSFER=1 but hf_transfer is not installed; disabling fast transfer."
|
||||
export HF_HUB_ENABLE_HF_TRANSFER=0
|
||||
fi
|
||||
else
|
||||
export HF_HUB_ENABLE_HF_TRANSFER=0
|
||||
fi
|
||||
|
||||
visible_gpus=$(nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null | wc -l | tr -d ' ')
|
||||
if [[ "$visible_gpus" -lt "$GPU_NUM" ]]; then
|
||||
echo "Expected at least $GPU_NUM GPUs, found $visible_gpus" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
for gpu_id in $(seq 0 $((GPU_NUM - 1))); do
|
||||
free_mb=$(nvidia-smi --id="$gpu_id" --query-gpu=memory.free --format=csv,noheader,nounits | tr -d ' ')
|
||||
if [[ "$free_mb" -lt "$MIN_FREE_GPU_MB" ]]; then
|
||||
echo "GPU $gpu_id has only ${free_mb} MiB free; text-only Wan preprocessing needs" \
|
||||
"about ${MIN_FREE_GPU_MB} MiB." >&2
|
||||
echo "Free the GPU or lower MIN_FREE_GPU_MB if you know this run will fit." >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
MARKER="$OUTPUT_DIR/.fastvideo_text_only_dmd_output"
|
||||
if [[ -d "$OUTPUT_DIR" && ! -f "$MARKER" ]]; then
|
||||
echo "Refusing to overwrite existing non-script output directory: $OUTPUT_DIR" >&2
|
||||
echo "Set OUTPUT_DIR to a new path or remove the directory manually." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
rm -rf "$OUTPUT_DIR"
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
touch "$MARKER"
|
||||
|
||||
echo "Text-only DMD preprocessing config:"
|
||||
echo " input: $INPUT_FILE"
|
||||
echo " output: $OUTPUT_DIR"
|
||||
echo " model: $MODEL_PATH"
|
||||
echo " gpus: $GPU_NUM"
|
||||
echo " batch size per GPU: $BATCH_SIZE"
|
||||
echo " text max length: $TEXT_MAX_LENGTH"
|
||||
echo " samples per parquet file: $SAMPLES_PER_FILE"
|
||||
echo " flush frequency: $FLUSH_FREQUENCY"
|
||||
|
||||
SHARD_DIR="$OUTPUT_DIR/_text_shards"
|
||||
mkdir -p "$SHARD_DIR"
|
||||
|
||||
python - "$INPUT_FILE" "$SHARD_DIR" "$GPU_NUM" <<'PY'
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
input_path = Path(sys.argv[1])
|
||||
shard_dir = Path(sys.argv[2])
|
||||
num_shards = int(sys.argv[3])
|
||||
|
||||
prompts = [line.rstrip("\n") for line in input_path.read_text(encoding="utf-8").splitlines() if line.strip()]
|
||||
if not prompts:
|
||||
raise SystemExit(f"No non-empty prompts found in {input_path}")
|
||||
|
||||
for shard_idx in range(num_shards):
|
||||
shard_prompts = prompts[shard_idx::num_shards]
|
||||
shard_path = shard_dir / f"train_text_shard_{shard_idx}.txt"
|
||||
shard_path.write_text("\n".join(shard_prompts) + "\n", encoding="utf-8")
|
||||
print(f"Wrote {len(shard_prompts)} prompts to {shard_path}")
|
||||
PY
|
||||
|
||||
run_preprocess_worker() {
|
||||
local gpu_id="$1"
|
||||
local shard_file="$2"
|
||||
local shard_output="$3"
|
||||
local log_file="$4"
|
||||
local master_port="$5"
|
||||
local -a cmd=(
|
||||
torchrun
|
||||
--nnodes=1
|
||||
--nproc_per_node=1
|
||||
--master_port "$master_port"
|
||||
fastvideo/pipelines/preprocess/v1_preprocess.py
|
||||
--model_path "$MODEL_PATH"
|
||||
--data_merge_path "$shard_file"
|
||||
--preprocess_video_batch_size "$BATCH_SIZE"
|
||||
--seed 42
|
||||
--max_height 448
|
||||
--max_width 832
|
||||
--num_frames 77
|
||||
--dataloader_num_workers 0
|
||||
--output_dir "$shard_output"
|
||||
--train_fps 16
|
||||
--samples_per_file "$SAMPLES_PER_FILE"
|
||||
--flush_frequency "$FLUSH_FREQUENCY"
|
||||
--text_max_length "$TEXT_MAX_LENGTH"
|
||||
--video_length_tolerance_range 5
|
||||
--preprocess_task text_only
|
||||
)
|
||||
|
||||
{
|
||||
echo "[gpu${gpu_id}] log file: $log_file"
|
||||
echo "[gpu${gpu_id}] command: CUDA_VISIBLE_DEVICES=${gpu_id} ${cmd[*]}"
|
||||
} | tee "$log_file"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$gpu_id" "${cmd[@]}" 2>&1 \
|
||||
| sed -u "s/^/[gpu${gpu_id}] /" \
|
||||
| tee -a "$log_file"
|
||||
local status=${PIPESTATUS[0]}
|
||||
if [[ "$status" -ne 0 ]]; then
|
||||
echo "[gpu${gpu_id}] preprocessing failed with exit code $status" | tee -a "$log_file"
|
||||
fi
|
||||
return "$status"
|
||||
}
|
||||
|
||||
pids=()
|
||||
for gpu_id in $(seq 0 $((GPU_NUM - 1))); do
|
||||
shard_file="$SHARD_DIR/train_text_shard_${gpu_id}.txt"
|
||||
shard_output="$OUTPUT_DIR/shard_${gpu_id}"
|
||||
mkdir -p "$shard_output"
|
||||
log_file="$OUTPUT_DIR/preprocess_gpu_${gpu_id}.log"
|
||||
|
||||
echo "Launching text-only preprocessing on GPU ${gpu_id}: ${shard_file}"
|
||||
run_preprocess_worker "$gpu_id" "$shard_file" "$shard_output" "$log_file" "$((29610 + gpu_id))" &
|
||||
pids+=("$!")
|
||||
done
|
||||
|
||||
failed=0
|
||||
for pid in "${pids[@]}"; do
|
||||
if ! wait "$pid"; then
|
||||
failed=1
|
||||
fi
|
||||
done
|
||||
|
||||
if [[ "$failed" -ne 0 ]]; then
|
||||
echo "One or more preprocessing workers failed. Check logs under $OUTPUT_DIR." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
num_parquet=$(find "$OUTPUT_DIR" -name '*.parquet' | wc -l | tr -d ' ')
|
||||
if [[ "$num_parquet" -eq 0 ]]; then
|
||||
echo "No parquet files were produced under $OUTPUT_DIR" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Text-only preprocessing complete."
|
||||
echo "Parquet files: $num_parquet"
|
||||
echo "Use this training data_path:"
|
||||
echo "$OUTPUT_DIR"
|
||||
Reference in New Issue
Block a user